From 33fe71a4d3316c847ccd44a35a2055a361148600 Mon Sep 17 00:00:00 2001 From: fxmarty <9808326+fxmarty@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:31:23 +0200 Subject: [PATCH 01/67] [AMD] Revert `Mxfp4MoeBackend.TRITON_UNFUSED` fallback (#46491) Signed-off-by: Felix Marty Co-authored-by: Felix Marty Co-authored-by: Andreas Karatzas --- tests/kernels/moe/test_ocp_mx_moe.py | 48 +++++++++++++++++++ tests/quantization/test_gfx950_moe.py | 15 ------ .../layers/fused_moe/oracle/mxfp4.py | 35 ++++++-------- 3 files changed, 63 insertions(+), 35 deletions(-) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 2f819c09aaa..7e8d4e3028a 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1558,3 +1558,51 @@ def test_mxfp4_emulation_rounds_up_to_block_size( # The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate. assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0 assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0 + + +def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( + monkeypatch: pytest.MonkeyPatch, +): + """ + select_mxfp4_moe_backend() must raise NotImplementedError, with the + collected per-backend unsupported reasons in the message, when no + backend supports the requested deployment configuration. + """ + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle + from vllm.model_executor.layers.fused_moe import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + RoutingMethodType, + ) + + class UnsupportedExperts: + @staticmethod + def is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ): + return False, f"unsupported reason for {cls.__name__}" + + monkeypatch.setattr( + mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [UnsupportedExperts] + ) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + monkeypatch.setattr(current_platform, "is_xpu", lambda: False) + monkeypatch.setattr(current_platform, "is_cpu", lambda: False) + + moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cpu", + routing_method=RoutingMethodType.Renormalize, + ) + + with pytest.raises(NotImplementedError, match="Unsupported reasons"): + mxfp4_oracle.select_mxfp4_moe_backend(moe_config) diff --git a/tests/quantization/test_gfx950_moe.py b/tests/quantization/test_gfx950_moe.py index 0efcc8a3c62..c8d34bb0ab5 100644 --- a/tests/quantization/test_gfx950_moe.py +++ b/tests/quantization/test_gfx950_moe.py @@ -79,21 +79,6 @@ def test_w4a4_dispatches_to_aiter(mxfp4_oracle_config): assert experts_cls is not None -@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") -@pytest.mark.skipif( - ROCM_AITER_AVAILABLE, - reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)", -) -def test_w4a4_falls_back_to_triton_unfused_without_aiter(mxfp4_oracle_config): - """Without AITER and no --moe-backend, ROCm falls back to TRITON_UNFUSED.""" - config = _make_w4a4_moe_config() - backend, experts_cls = select_mxfp4_moe_backend( - config, activation_key=kMxfp4Dynamic - ) - assert backend == Mxfp4MoeBackend.TRITON_UNFUSED - assert experts_cls is not None - - @pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)") def test_w4a4_dispatches_to_emulation_with_moe_backend(mxfp4_oracle_config): """With --moe-backend emulation, W4A4 selects EMULATION.""" diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 921e8f114d5..07cff8fa22c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -502,6 +502,7 @@ def select_mxfp4_moe_backend( _get_priority_backends_for_gpt_oss(), requested_activation_key ) + unsupported_reasons = [] for backend in AVAILABLE_BACKENDS: # Use requested_activation_key if provided, otherwise use backend default act_key = ( @@ -518,6 +519,7 @@ def select_mxfp4_moe_backend( return backend, k_cls else: logger.debug_once(_make_log_unsupported(backend, reason)) + unsupported_reasons.append((backend, reason)) if current_platform.is_xpu(): backend = Mxfp4MoeBackend.XPU @@ -541,26 +543,19 @@ def select_mxfp4_moe_backend( activation_format, ) - if current_platform.is_rocm(): - backend = Mxfp4MoeBackend.TRITON_UNFUSED - logger.info_once(_make_log_backend(backend)) - return _return_or_raise( - Mxfp4MoeBackend.TRITON_UNFUSED, - config, - kMxfp4Static, - None, - activation_format, - ) - - if current_platform.is_cuda(): - raise NotImplementedError( - "No MXFP4 MoE backend supports the deployment configuration. " - f"weight_key=kMxfp4Static, activation_key={activation_key}. " - "Native backends require specific hardware. " - "Set `VLLM_LOGGING_LEVEL=DEBUG` to see detailed unsupported reasons. " - ) - - return Mxfp4MoeBackend.NONE, None + unsupported_log = "; ".join( + [ + f"backend: {backend.value}, reason: {reason}" + for backend, reason in unsupported_reasons + ] + ) + raise NotImplementedError( + "No MXFP4 MoE backend supports the deployment configuration. " + f"weight_key=kMxfp4Static, activation_key={activation_key}. " + f"Candidate backends were: " + f"{[backend.value for backend in AVAILABLE_BACKENDS]}. " + f"Unsupported reasons: {unsupported_log}. " + ) def select_deepseek_v4_mxfp4_moe_backend( From fbb1ef68030991a14291231bf8877063b57e1ada Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:32:42 -0700 Subject: [PATCH 02/67] [Bugfix] Fix DeepseekV4FP8 Quark MXFP4 crash on list-valued weight (#49634) Signed-off-by: Colin Zeng Co-authored-by: Andreas Karatzas --- vllm/models/deepseek_v4/quant_config.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index 89cf695baf0..293d71f2f41 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -120,7 +120,11 @@ class DeepseekV4FP8Config(Fp8Config): @staticmethod def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool: """True for AMD-Quark exports whose global scheme is MXFP4.""" - weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") or {} + weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") + # A non-dict weight (e.g. a list of multiple specs) means not an OCP + # MXFP4 scheme (e.g. NVFP4 with 2-level scale). + if not isinstance(weight, dict): + return False return ( weight.get("dtype") == "fp4" and weight.get("qscheme") == "per_group" From a8f296083f43e236f908d4d898ccfb0b108db1f6 Mon Sep 17 00:00:00 2001 From: Chang Guo Date: Mon, 27 Jul 2026 21:10:20 -0700 Subject: [PATCH 03/67] [KV Offload] Make compact secondary identity TP-independent (#49858) Signed-off-by: Change72 Co-authored-by: GPT-5.6 Sol Co-authored-by: Or Ozeri --- tests/v1/kv_offload/test_file_mapper.py | 105 +++++++++++++++++++ tests/v1/kv_offload/tiering/test_fs_tier.py | 66 +++++++++++- tests/v1/kv_offload/tiering/test_obj_tier.py | 55 +++++++++- vllm/v1/kv_offload/file_mapper.py | 11 +- 4 files changed, 226 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index c2c4427e184..027f01133b2 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -51,6 +51,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: data_parallel_index=0, is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), ), + replicated_layout=kwargs.get("replicated_layout", False), ) spec = MagicMock(spec=OffloadingSpec) spec.config = config @@ -205,3 +206,107 @@ def test_parallel_agnostic_separates_persistent_layouts(): assert agnostic.base_path != specific.base_path assert "parallel_agnostic" not in agnostic.fields assert specific.fields["parallel_agnostic"] is False + + +# --------------------------------------------------------------------------- +# replicated_layout: OR'd into parallel-agnostic identity for compact rows +# --------------------------------------------------------------------------- + + +def test_replicated_layout_collapses_parallel_identity(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=1, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=3, **shared) + + assert tp2.base_path == tp4.base_path + for fm in (tp2, tp4): + assert fm.fields["tp_size"] == 1 + assert fm.fields["pp_size"] == 1 + assert fm.fields["pcp_size"] == 1 + assert fm.fields["dcp_size"] == 1 + assert fm.rank == 0 + assert "parallel_agnostic" not in fm.fields + assert fm.fields["replicated_layout"] is True + assert fm.get_file_name(make_offload_key(b"\x01" * 8, 0)).startswith( + f"{fm.base_path}_r0/" + ) + + +def test_replicated_layout_requires_caller_opt_in(): + fm = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=True, + parallel_agnostic=False, + ) + assert fm.fields["tp_size"] == 2 + assert fm.rank == 1 + assert fm.fields["parallel_agnostic"] is False + assert "replicated_layout" not in fm.fields + baseline = make_mapper_from_offloading_spec( + tp_size=2, + world_size=2, + rank=1, + replicated_layout=False, + parallel_agnostic=False, + ) + assert fm.base_path == baseline.base_path + + +def test_non_replicated_keeps_parallel_identity(): + fm = make_mapper_from_offloading_spec( + tp_size=4, + world_size=4, + rank=2, + replicated_layout=False, + is_parallelism_agnostic=False, + parallel_agnostic=True, + ) + assert fm.fields["tp_size"] == 4 + assert fm.rank == 2 + assert fm.fields["parallel_agnostic"] is False + assert fm.get_file_name(make_offload_key(b"\x02" * 8, 0)).startswith( + f"{fm.base_path}_r2/" + ) + + +def test_replicated_and_parallelism_agnostic_separate_layouts(): + shared = dict( + model_name="shared-model", + groups=((16, "layer0"),), + tp_size=2, + world_size=2, + rank=1, + parallel_agnostic=True, + ) + via_agnostic = make_mapper_from_offloading_spec( + is_parallelism_agnostic=True, + replicated_layout=False, + **shared, + ) + via_replicated = make_mapper_from_offloading_spec( + is_parallelism_agnostic=False, + replicated_layout=True, + **shared, + ) + assert via_agnostic.base_path != via_replicated.base_path + assert "replicated_layout" not in via_agnostic.fields + assert via_replicated.fields["replicated_layout"] is True + + +def test_replicated_layout_run_config_tp_invariant(): + shared = dict( + model_name="mla-model", + groups=((16, "mla_layer"),), + replicated_layout=True, + parallel_agnostic=True, + ) + tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared) + tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared) + assert tp2.get_run_config() == tp4.get_run_config() diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 66533e3676b..2959ac1aa03 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -52,8 +52,18 @@ _DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") -def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: +def _make_offloading_spec( + enable_kv_cache_events: bool = False, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> MagicMock: """Mock spec with an explicit global KV events flag.""" + if world_size is None: + world_size = tp_size spec = MagicMock() spec.config = OffloadingConfig( groups=(), @@ -64,15 +74,16 @@ def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: model=OffloadingModelConfig(name="test-model", dtype="float32"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) spec.blocks_per_chunk = 1 spec.kv_events_config = OffloadingKVEventsConfig( @@ -725,3 +736,48 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): assert not fs_events[0].removed finally: tier.shutdown() + + +def test_fs_tier_cross_tp_round_trip(tmp_path): + """TP=2 replicated writer and TP=4 reader share namespace and bytes.""" + root = str(tmp_path) + writer_tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) + expected = writer_tensor[0].clone() + writer = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=2, world_size=2, rank=0, replicated_layout=True + ), + primary_kv_view=memoryview(writer_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + writer.submit_store(make_job(1, [key(7)], [0])) + assert all(r.success for r in drain(writer)) + writer_base = writer.file_mapper.base_path + writer_path = writer.file_mapper.get_file_name(key(7)) + finally: + writer.shutdown() + + reader_tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + reader = FileSystemTierManager( + offloading_spec=_make_offloading_spec( + tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + primary_kv_view=memoryview(reader_tensor.numpy()), + tier_type="fs", + root_dir=root, + n_read_threads=2, + n_write_threads=2, + ) + try: + assert reader.file_mapper.base_path == writer_base + assert reader.file_mapper.get_file_name(key(7)) == writer_path + assert lookup_and_wait(reader, [key(7)]) == [LookupResult.HIT] + reader.submit_load(make_job(2, [key(7)], [1], is_promotion=True)) + assert all(r.success for r in drain(reader)) + assert torch.allclose(reader_tensor[1], expected) + finally: + reader.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index fc30e1437a7..82ba183ea5f 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -43,7 +43,17 @@ from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManag # --------------------------------------------------------------------------- -def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: +def _make_offloading_config( + enable_kv_cache_events: bool, + *, + tp_size: int = 1, + rank: int = 0, + world_size: int | None = None, + replicated_layout: bool = False, + is_parallelism_agnostic: bool = False, +) -> OffloadingConfig: + if world_size is None: + world_size = tp_size return OffloadingConfig( groups=(), worker_kv_bytes_per_block=0, @@ -53,15 +63,16 @@ def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: model=OffloadingModelConfig(name="test/model", dtype="float16"), cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), parallel=OffloadingParallelConfig( - rank=0, - world_size=1, - tp_size=1, + rank=rank, + world_size=world_size, + tp_size=tp_size, pp_size=1, pcp_size=1, dcp_size=1, data_parallel_index=0, - is_parallelism_agnostic=False, + is_parallelism_agnostic=is_parallelism_agnostic, ), + replicated_layout=replicated_layout, ) @@ -617,3 +628,37 @@ class TestObjStoreConfig: params = cfg.to_nixl_params() assert params["ca_bundle"] == "/path/to/ca.pem" assert "access_key" not in params + + +def test_obj_tier_replicated_layout_collapses_mapper_identity(): + """TP=2 and TP=4 replicated configs share the obj FileMapper namespace.""" + tp2_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=2, world_size=2, rank=1, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp4_spec = SimpleNamespace( + config=_make_offloading_config( + False, tp_size=4, world_size=4, rank=3, replicated_layout=True + ), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, + self_describing_kv_events=False, + ), + ) + tp2_tier, _ = _make_tier(offloading_spec=tp2_spec) + tp4_tier, _ = _make_tier(offloading_spec=tp4_spec) + try: + assert tp2_tier._file_mapper.base_path == tp4_tier._file_mapper.base_path + assert tp2_tier._file_mapper.rank == 0 + assert tp4_tier._file_mapper.rank == 0 + assert tp2_tier._file_mapper.get_run_config() == ( + tp4_tier._file_mapper.get_run_config() + ) + finally: + tp2_tier.shutdown() + tp4_tier.shutdown() diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index 8e8c19d53d6..4b12dba913d 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -35,6 +35,7 @@ class FileMapper: kv_cache_groups: list[dict] | None = None, inference_engine: str = "vllm", parallel_agnostic: bool = False, + replicated_layout: bool = False, ): """ Initialize the file mapper. Each worker constructs its own, but @@ -60,6 +61,10 @@ class FileMapper: } if not parallel_agnostic: self.fields["parallel_agnostic"] = False + # Only written when True so existing deployments' hashed fields are + # unchanged (False is the historical default and must not appear). + if replicated_layout: + self.fields["replicated_layout"] = True self.base_path: str = self._compute_base_path(root_dir, self.fields) @classmethod @@ -92,7 +97,11 @@ class FileMapper: rank=parallel.rank, dtype=config.model.dtype, kv_cache_groups=kv_cache_groups, - parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic), + parallel_agnostic=( + parallel_agnostic + and (parallel.is_parallelism_agnostic or config.replicated_layout) + ), + replicated_layout=(parallel_agnostic and config.replicated_layout), ) def get_file_name(self, key: OffloadKey) -> str: From 52c3c4a42fd13b62ba985b9ceb9b9969964ee83e Mon Sep 17 00:00:00 2001 From: MINJUN GIL Date: Tue, 28 Jul 2026 13:10:52 +0900 Subject: [PATCH 04/67] [Bugfix][KV Offload][OBJ] Preserve job completion during cleanup (#49947) Signed-off-by: MINJUN GIL Co-authored-by: OpenAI Codex --- tests/v1/kv_offload/tiering/test_obj_tier.py | 111 ++++++++++++++++++- vllm/v1/kv_offload/tiering/obj/manager.py | 29 ++++- 2 files changed, 133 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 82ba183ea5f..661438dce63 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -35,6 +35,10 @@ from vllm.v1.kv_offload.config import ( OffloadingParallelConfig, ) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.manager import ( + CPUPrimaryTierOffloadingManager, + TieringOffloadingManager, +) from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -220,12 +224,14 @@ def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: def _make_tier( num_blocks: int = 4, offloading_spec: SimpleNamespace = _OFFLOADING_SPEC, + primary_kv_view: memoryview | None = None, **tier_kwargs, ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: """Create a tier backed by a fresh MockNixlAgent.""" mock_agent = MockNixlAgent() - tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) - view = memoryview(tensor.numpy()) + if primary_kv_view is None: + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) with ( patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"), patch( @@ -235,7 +241,7 @@ def _make_tier( ): tier = ObjectStoreSecondaryTierManager( offloading_spec=offloading_spec, - primary_kv_view=view, + primary_kv_view=primary_kv_view, tier_type="obj", store_config=_STORE_CONFIG, prefix=_RUN_PREFIX, @@ -449,6 +455,105 @@ class TestMockObjTierFailures: assert not by_id[1].success assert by_id[2].success + def test_release_xfer_failure_retries_without_losing_result(self, monkeypatch): + tier, agent = _make_tier(num_blocks=4) + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + + tier.submit_store(make_job(1, [key(1)], [0])) + + # The transfer handle could not be released safely, so the job must + # remain tracked and must not be finalized yet. + assert list(tier.get_finished_jobs()) == [] + assert 1 in tier._transfers + + # Cleanup is retried without polling again or changing the failure + # verdict. The completion is then returned exactly once. + results = list(tier.get_finished_jobs()) + assert len(results) == 1 + assert results[0].job_id == 1 + assert not results[0].success + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + @pytest.mark.parametrize( + "cleanup_method", ["release_dlist_handle", "deregister_memory"] + ) + def test_post_transfer_cleanup_failure_does_not_lose_result( + self, monkeypatch, cleanup_method + ): + tier, agent = _make_tier(num_blocks=4) + monkeypatch.setattr( + agent, + cleanup_method, + MagicMock(side_effect=RuntimeError("cleanup failed")), + ) + + tier.submit_store(make_job(1, [key(1)], [0])) + results = list(tier.get_finished_jobs()) + + assert len(results) == 1 + assert results[0].job_id == 1 + assert results[0].success + assert not tier._transfers + assert list(tier.get_finished_jobs()) == [] + + def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch): + num_blocks = 4 + tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) + primary_kv_view = memoryview(tensor.numpy()) + mmap_region = MagicMock() + mmap_region.create_kv_memoryview.return_value = primary_kv_view + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=num_blocks, mmap_region=mmap_region + ) + obj_tier, agent = _make_tier( + num_blocks=num_blocks, primary_kv_view=primary_kv_view + ) + manager = TieringOffloadingManager( + primary_tier=primary_tier, secondary_tiers=[obj_tier] + ) + + keys = [key(1)] + primary_result = primary_tier.prepare_store(keys, _CTX) + assert primary_result is not None + primary_tier.complete_store(keys, _CTX, success=True) + job = manager.create_store_job(keys, _CTX) + obj_tier.submit_store(job) + + block = primary_tier._policy.get(keys[0]) + assert block is not None + assert block.ref_cnt == 1 + assert len(manager._transfer_jobs) == 1 + + agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) + release_xfer = MagicMock( + side_effect=[RuntimeError("transfer is still active"), None] + ) + monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) + schedule_context = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + + manager.on_schedule_end(schedule_context) + + assert len(manager._transfer_jobs) == 1 + assert block.ref_cnt == 1 + assert len(obj_tier._transfers) == 1 + assert manager.has_pending_work() + + manager.on_schedule_end(schedule_context) + + assert manager._transfer_jobs == {} + assert block.ref_cnt == 0 + assert obj_tier._transfers == {} + assert not manager.has_pending_work() + assert agent.check_xfer_state.call_count == 2 + assert release_xfer.call_count == 2 + class TestMockObjTierShutdown: def test_shutdown_clears_in_flight_transfers(self): diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index c7e0d4c4beb..2dfc2d30fa2 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -307,15 +307,36 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): else: if state == NIXL_PROC: continue - elif state == NIXL_DONE: + if state == NIXL_DONE: success = True else: success = False logger.warning("transfer failed job=%d state=%s", job_id, state) + + try: + self._agent.release_xfer_handle(entry.xfer_handle) + except Exception as exc: + # Keep the entry until NIXL confirms that the transfer handle + # can be released. The transfer may still access primary-tier + # memory, so publishing its result would allow unsafe reuse. + logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc) + continue + + # Once the transfer handle is released, these remaining cleanup + # failures must not suppress the job completion. They can leak + # NIXL metadata, but cannot leave an active data transfer behind. + try: + self._agent.release_dlist_handle(entry.obj_handle) + except Exception as exc: + logger.warning( + "release_dlist_handle failed for job %d: %s", job_id, exc + ) + try: + self._agent.deregister_memory(entry.files_desc) + except Exception as exc: + logger.warning("deregister_memory failed for job %d: %s", job_id, exc) + del self._transfers[job_id] - self._agent.release_xfer_handle(entry.xfer_handle) - self._agent.release_dlist_handle(entry.obj_handle) - self._agent.deregister_memory(entry.files_desc) self._pending_results.append(JobResult(job_id=job_id, success=success)) def get_finished_jobs(self) -> Iterable[JobResult]: From d223c900d85224c02f2162ee2c757a769e99f519 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 27 Jul 2026 21:36:45 -0700 Subject: [PATCH 05/67] [Bugfix] Only pad transformers backend `value` when it is narrower (#50060) Signed-off-by: Nick Hill --- vllm/model_executor/models/transformers/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 9dbfe4b5031..ff4bb6cd433 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -64,12 +64,15 @@ def vllm_attention_forward( head_dim_v = value.shape[-1] query, key, value = (x.transpose(1, 2) for x in (query, key, value)) query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) - # Pad `value` up to the query/key head size when they differ (expanded MLA). - if head_dim_v != head_dim_qk: + # Pad `value` up to the query/key head size when it is smaller (expanded + # MLA). A larger last dim just means `value` isn't split per head, e.g. + # packed grouped/multi-query projections, and needs no padding. + pad_value = head_dim_v < head_dim_qk + if pad_value: value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v)) value = value.reshape(hidden, -1) attn_output = self_attn.forward(query, key, value) - if head_dim_v != head_dim_qk: + if pad_value: attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v] attn_output = attn_output.reshape(hidden, -1) return attn_output, None From 74587939b17b4eb3ba281ec9ca43f93aa230a3cb Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Tue, 28 Jul 2026 01:01:56 -0400 Subject: [PATCH 06/67] [Build] Fix CUDA arch detection producing kernel-less builds on SM121 (#49904) --- CMakeLists.txt | 13 +++++++++---- cmake/utils.cmake | 11 ++++++----- tests/test_cmake_utils.py | 25 +++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 9 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9bd73f18e70..3a4e23ad0b3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -219,10 +219,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # the set of architectures we want to compile for and remove the from the # CMAKE_CUDA_FLAGS so that they are not applied globally. # - # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. It is emitted by torch - # as `code=compute_*`, while extract_unique_cuda_archs_ascending() records only - # `arch=compute_*`. If a kernel really needs PTX, add `+PTX` to that kernel's - # component-specific arch list below. + # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really + # needs PTX, add `+PTX` to that kernel's component-specific arch list below. # clear_cuda_arches(CUDA_ARCH_FLAGS) extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") @@ -232,6 +230,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") cuda_archs_loose_intersection(CUDA_ARCHS "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}") message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}") + if(NOT CUDA_ARCHS) + message(FATAL_ERROR + "No supported CUDA architectures; the build would produce a binary " + "with no usable kernels. Detected gencode flags: ${CUDA_ARCH_FLAGS}; " + "supported: ${CUDA_SUPPORTED_ARCHS}. " + "Set TORCH_CUDA_ARCH_LIST for your GPU (e.g. 12.0).") + endif() else() # # For other GPU targets override the GPU architectures detected by cmake/torch diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 14a94eebb22..bbae89c1f57 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -241,14 +241,15 @@ endmacro() # `.`, dedupes them and then sorts them in ascending order and # stores them in `OUT_ARCHES`. # -# Example: -# CUDA_ARCH_FLAGS="-gencode arch=compute_75,code=sm_75;...;-gencode arch=compute_90a,code=sm_90a" -# extract_unique_cuda_archs_ascending(OUT_ARCHES CUDA_ARCH_FLAGS) -# OUT_ARCHES="7.5;...;9.0" +# Prefer `code=sm_*`; fall back to `arch=compute_*` for PTX-only flags. +# This handles mismatches such as `arch=compute_20,code=sm_121`. function(extract_unique_cuda_archs_ascending OUT_ARCHES CUDA_ARCH_FLAGS) set(_CUDA_ARCHES) foreach(_ARCH ${CUDA_ARCH_FLAGS}) - string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + string(REGEX MATCH "code=sm_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + if (NOT _COMPUTE) + string(REGEX MATCH "arch=compute_\([0-9]+[af]?\)" _COMPUTE ${_ARCH}) + endif() if (_COMPUTE) set(_COMPUTE ${CMAKE_MATCH_1}) endif() diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index 227ec231eb2..d0673bc462e 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -21,3 +21,28 @@ endif() ) subprocess.run(["cmake", "-P", script], check=True) + + +def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( + tmp_path: Path, +): + """torch's autodetection can emit a bogus arch=compute_* half (e.g. + capability 12.1 corrupted to arch=compute_20,code=sm_121); the SASS + target must win, while PTX-only entries keep the virtual arch.""" + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_extract_archs.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +extract_unique_cuda_archs_ascending(actual + "-gencode arch=compute_20,code=sm_121;\ +-gencode arch=compute_80,code=sm_80;\ +-gencode arch=compute_80,code=compute_80") +if(NOT "${{actual}}" STREQUAL "8.0;12.1") + message(FATAL_ERROR "Expected '8.0;12.1', got '${{actual}}'") +endif() +""" + ) + + subprocess.run(["cmake", "-P", script], check=True) From f472ab0a4c0987d2badc2f2aeecd38fee4180a4c Mon Sep 17 00:00:00 2001 From: afriedri Date: Tue, 28 Jul 2026 00:53:46 -0500 Subject: [PATCH 07/67] Remove triton per group quant [ROCm] [Bugfix] (#49621) Signed-off-by: Andy Friedrich --- .../distributed/test_fusion_all_reduce.py | 31 ++++------------- tests/compile/passes/test_fusion.py | 2 -- .../passes/test_silu_mul_quant_fusion.py | 11 +----- .../passes/fusion/allreduce_rms_fusion.py | 3 +- .../layers/quantization/input_quant_fp8.py | 5 --- .../layers/quantization/utils/fp8_utils.py | 34 ------------------- 6 files changed, 8 insertions(+), 78 deletions(-) diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1aac4b2bec4..e9c8d0deaa7 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -272,12 +272,10 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): token_num=16, eps=1e-6, dtype: torch.dtype = torch.bfloat16, - use_triton_quant: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps - self.use_triton_quant = use_triton_quant assert hidden_size % self.quant_group_size == 0, ( f"hidden_size ({hidden_size}) must be a multiple of " f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant" @@ -289,10 +287,6 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): ] def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - if self.use_triton_quant: - return torch.ops.vllm.triton_per_token_group_quant_fp8( - rms, self.quant_group_size - ) return torch.ops.vllm.rocm_aiter_group_fp8_quant.default( rms, self.quant_group_size ) @@ -339,11 +333,7 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): def ops_in_model_before(self): return [ torch.ops.vllm.all_reduce.default, - ( - torch.ops.vllm.triton_per_token_group_quant_fp8.default - if self.use_triton_quant - else torch.ops.vllm.rocm_aiter_group_fp8_quant.default - ), + torch.ops.vllm.rocm_aiter_group_fp8_quant.default, ] def ops_in_model_after(self): @@ -646,7 +636,6 @@ def all_reduce_fusion_pass_on_test_model( @multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize("use_triton_quant", [True, False]) @pytest.mark.parametrize("batch_size", [8]) @pytest.mark.parametrize("seq_len", [8]) @pytest.mark.parametrize("hidden_size", [128]) @@ -663,7 +652,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -676,9 +664,9 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( * ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual, single ``rms`` consumer) * ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with- - residual, DSv3.2 indexer fan-out; parametrized over both - ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant`` - producers). + residual, DSv3.2 indexer fan-out; parametrized over + ``rocm_aiter_group_fp8_quant`` + producer). """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -703,7 +691,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size, dtype, enable_rms_norm_custom_op, - use_triton_quant, monkeypatch, ), nprocs=nprocs, @@ -721,7 +708,6 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): set_random_seed(0) @@ -749,10 +735,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( custom_ops = [] if enable_rms_norm_custom_op: custom_ops.append("+rms_norm") - # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip`` - # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at - # the call site). The patterns in this PR are robust to both Triton and - # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example + # We always enable +quant_fp8 so the matcher's example # trace finds the same form the test model uses. custom_ops.append("+quant_fp8") @@ -783,9 +766,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - model = test_model_cls( - hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant - ) + model = test_model_cls(hidden_size, token_num, dtype=dtype) hidden_states = torch.randn((token_num, hidden_size), requires_grad=False) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 92d1902b2c2..591b014d9e2 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -195,8 +195,6 @@ class TestModel(torch.nn.Module): # Blockwise path if self.use_aiter_fusion and self.use_aiter_quant_op: return [rocm_aiter_ops.get_group_quant_op()] - if self.use_aiter_fusion: - return [torch.ops.vllm.triton_per_token_group_quant_fp8.default] else: if self.use_aiter_quant_op: return [rocm_aiter_ops.get_per_token_quant_op()] diff --git a/tests/compile/passes/test_silu_mul_quant_fusion.py b/tests/compile/passes/test_silu_mul_quant_fusion.py index bc134ed427a..7d291cc5044 100644 --- a/tests/compile/passes/test_silu_mul_quant_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_fusion.py @@ -158,13 +158,6 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): input_dtype=dtype, ) - if not current_platform.is_fp8_fnuz(): - kernel = self.w8a8_block_fp8_linear.kernel - orig_quant = kernel.quant_fp8 - kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant( - *a, use_triton=True, **kw - ) - self.enable_silu_mul_custom_op = self.silu_and_mul.enabled() def forward(self, x): @@ -175,9 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): def ops_in_model_before(self): return [ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul, - rocm_aiter_ops.get_group_quant_op() - if current_platform.is_fp8_fnuz() - else torch.ops.vllm.triton_per_token_group_quant_fp8.default, + rocm_aiter_ops.get_group_quant_op(), ] def ops_in_model_after(self): diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index f7ff7df66cd..1722b524eeb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1416,8 +1416,7 @@ class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent with the sibling patterns above), which traces both ``QuantFP8.forward_hip`` and ``forward_native`` paths and so matches whichever op the call site - lowers to (``vllm.triton_per_token_group_quant_fp8`` or - ``vllm.rocm_aiter_group_fp8_quant``). + lowers to (``vllm.rocm_aiter_group_fp8_quant``). """ def __init__( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index e8810919c20..2eb34630aa6 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -139,11 +139,6 @@ class QuantFP8(CustomOp): scale_ub: torch.Tensor | None = None, use_triton: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if self.is_group_quant and use_triton: - assert scale is None, "Dynamic group quantization does not use scale" - - return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size) - use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous() use_aiter_per_tensor_quant = ( use_aiter_quant and self.group_shape.is_per_tensor() diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 83e56a4567b..2e4fbdf4c64 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -34,7 +34,6 @@ from vllm.utils.deep_gemm import ( transform_sf_into_required_layout, ) from vllm.utils.platform_utils import get_device_name_as_file_name -from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -45,39 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool: return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz -def _triton_per_token_group_quant_fp8_impl( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return per_token_group_quant_fp8( - x, group_size, column_major_scales=False, use_ue8m0=False - ) - - -def _triton_per_token_group_quant_fp8_fake( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - M, N = x.shape - x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device) - out_bs = torch.empty( - ( - M, - (N + group_size - 1) // group_size, - ), - dtype=torch.float32, - device=x.device, - ) - return x_fp8, out_bs - - -direct_register_custom_op( - "triton_per_token_group_quant_fp8", - _triton_per_token_group_quant_fp8_impl, - fake_impl=_triton_per_token_group_quant_fp8_fake, -) - - def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: From 90245f4190a35593a625e4bc349485c39c774d39 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Tue, 28 Jul 2026 14:09:48 +0800 Subject: [PATCH 08/67] [Bugfix] Fix multi-modal support on CPU MRV2 (#50073) Signed-off-by: jiang1.li --- vllm/multimodal/inputs.py | 2 ++ vllm/v1/worker/cpu/shm.py | 12 ++++++++++++ vllm/v1/worker/gpu/model_states/encoder_decoder.py | 5 ++++- 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/vllm/multimodal/inputs.py b/vllm/multimodal/inputs.py index 71f17a8648b..dc96c366ad2 100644 --- a/vllm/multimodal/inputs.py +++ b/vllm/multimodal/inputs.py @@ -457,6 +457,8 @@ class BaseMultiModalField(ABC): device = "cpu" if pin_memory and self.keep_on_cpu: pin_memory = False + if device == "cpu" or device == torch.device("cpu"): + pin_memory = False batch = [elem.data for elem in elems] out = self._reduce_data(batch, pin_memory=pin_memory) diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index bd1f96c71ed..970e8a2d414 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -24,12 +24,17 @@ def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tens class _EventPlaceholder: def __init__(self, *args, **kwargs) -> None: self.record = noop + self.wait = noop self.synchronize = noop class _StreamPlaceholder: def __init__(self, *args, **kwargs) -> None: self.wait_stream = noop + self.wait_event = noop + self.record_event = noop + self.synchronize = noop + self.query = lambda: True self.device = torch.device("cpu") def __enter__(self, *args, **kwargs): @@ -55,6 +60,7 @@ torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder() torch.accelerator.synchronize = noop torch.accelerator.empty_cache = noop torch.Tensor.pin_memory = fake_pin_memory +torch.Tensor.record_stream = noop torch.accelerator.get_memory_info = get_memory_info # Patch vLLM torch utils @@ -80,3 +86,9 @@ import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer + +# Patch Triton +from vllm.triton_utils import HAS_TRITON, tl + +if HAS_TRITON: + tl.debug_barrier = noop diff --git a/vllm/v1/worker/gpu/model_states/encoder_decoder.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py index f759c0b1e15..618984c97f3 100644 --- a/vllm/v1/worker/gpu/model_states/encoder_decoder.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -9,6 +9,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata from vllm.v1.worker.gpu.input_batch import InputBatch @@ -157,7 +158,9 @@ class EncoderDecoderModelState(ModelState): for_capture: bool, num_reqs: int, ) -> dict[int, tuple[torch.Tensor, np.ndarray]]: - encoder_seq_lens = torch.zeros(num_reqs, dtype=torch.int32, pin_memory=True) + encoder_seq_lens = torch.zeros( + num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY + ) encoder_seq_lens_np = encoder_seq_lens.numpy() if not for_capture: # During normal execution, use actual encoder lengths. From 9069a57139bc075733ffb82c228d9f33945e524b Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:35:29 +0800 Subject: [PATCH 09/67] [Core][Frontend] Add weight version tagging for RL rollouts (#49040) Signed-off-by: Shuolei Wang Signed-off-by: Shuolei Wang <948904026@qq.com> --- docs/serving/offline_inference.md | 2 ++ docs/serving/online_serving/README.md | 2 ++ docs/training/async_rl.md | 3 +- docs/training/weight_transfer/README.md | 6 ++-- tests/distributed/test_weight_transfer.py | 9 +++++- .../entrypoints/openai/test_openai_schema.py | 1 + .../test_weight_transfer_llm.py | 12 +++++++- vllm/distributed/weight_transfer/base.py | 2 +- vllm/distributed/weight_transfer/clients.py | 13 ++++++-- vllm/engine/protocol.py | 12 ++++++-- vllm/entrypoints/llm.py | 14 +++++++-- vllm/entrypoints/serve/dev/rlhf/api_router.py | 24 +++++++++++++-- vllm/v1/engine/async_llm.py | 14 +++++++-- vllm/v1/engine/core.py | 9 ++++++ vllm/v1/engine/core_client.py | 30 +++++++++++++++++++ vllm/v1/engine/llm_engine.py | 7 +++++ 16 files changed, 142 insertions(+), 18 deletions(-) diff --git a/docs/serving/offline_inference.md b/docs/serving/offline_inference.md index 4512f4a0720..9a71612f262 100644 --- a/docs/serving/offline_inference.md +++ b/docs/serving/offline_inference.md @@ -65,6 +65,8 @@ For further details on Weight Transfer, please refer to [this page](../training/ - `LLM.start_weight_update` - Starts a new weight update cycle. - `LLM.update_weights` - Updates the model weights. - `LLM.finish_weight_update` - Finishes the current weight update cycle. +- `LLM.update_weight_version` - Sets the weight version without updating model weights. +- `LLM.get_weight_version` - Returns the latest committed weight version. ## Additional APIs diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index f7914bc0582..90ff7a3e3d8 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -179,6 +179,8 @@ For further details on Weight Transfer, please refer to [this page](../../traini - `/start_weight_update` - Prepares the inference engine for a weight update. - `/update_weights` - Update model weights (can alter model behavior) - `/finish_weight_update` - Finalizes the weight update +- `/update_weight_version` - Set the weight version without updating model weights +- `/weight_info` - Get the latest committed weight version - `/get_world_size` - Get distributed world size ### Collective RPC diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index e655f9c39ff..9e75a24eaa1 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -38,11 +38,12 @@ Resumes the scheduler after a pause. Any requests frozen with `mode="keep"` will ### HTTP Endpoints -When using the vLLM HTTP server, the same functionality is available via: +With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionality via: - `POST /pause?mode=keep` - Pause generation - `POST /resume` - Resume generation - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`) +- `GET /weight_info` - Return the latest committed `weight_version` !!! note "Data Parallelism" When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update. diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 7579e5fd4d0..b8d39763181 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -53,7 +53,9 @@ When running vLLM as an HTTP server, the following endpoints are available for w | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info | | `/start_weight_update` | POST | Start a weight update | | `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata | -| `/finish_weight_update` | POST | Finish the weight update and run post-processing | +| `/finish_weight_update` | POST | Finish the update and optionally commit its `weight_version` | +| `/update_weight_version` | POST | Update `weight_version` without changing model weights | +| `/weight_info` | GET | Get the latest committed weight version | | `/pause` | POST | Pause generation before weight sync to handle inflight requests | | `/resume` | POST | Resume generation after weight sync | | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) | @@ -79,7 +81,7 @@ EngineClass.trainer_send_weights( ) # 4. Finish weight update on inference side -llm.finish_weight_update() +llm.finish_weight_update(weight_version="step-42") ``` See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples. diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index b79aa1974d1..eeeceb95998 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -1247,7 +1247,7 @@ class RecordingClient: self.order.append("update") self.last_update_info = update_info - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: self.order.append("finish") @@ -1303,6 +1303,10 @@ class TestTrainerClients: assert isinstance(update_req, WeightTransferUpdateRequest) assert update_req.update_info == {"names": ["w"]} + client.finish_weight_update("step-42") + handle.finish_weight_update.remote.assert_called_once_with() + handle.update_weight_version.remote.assert_called_once_with("step-42") + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" captured = {} @@ -1334,6 +1338,9 @@ class TestTrainerClients: client.update_weights(update_info) assert captured["json"]["update_info"] == update_info + client.finish_weight_update("step-42") + assert captured["json"] == {"weight_version": "step-42"} + class TestModuleSource: """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" diff --git a/tests/entrypoints/openai/test_openai_schema.py b/tests/entrypoints/openai/test_openai_schema.py index 2985c539518..6d3fc2f4474 100644 --- a/tests/entrypoints/openai/test_openai_schema.py +++ b/tests/entrypoints/openai/test_openai_schema.py @@ -148,6 +148,7 @@ def test_openapi_stateless(case: schemathesis.Case): "/start_draft_weight_update", "/update_weights", "/finish_weight_update", + "/update_weight_version", ): return diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 9088b3c5e8d..31d562e5ccf 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -234,6 +234,7 @@ def test_update_weights_calls_engine(): assert shapes == test_shapes llm.finish_weight_update() + assert llm.get_weight_version() == "default" @create_new_process_for_each_test() @@ -259,6 +260,8 @@ def test_full_weight_transfer_flow(): weight_transfer_config=WeightTransferConfig(backend="nccl"), ) + assert llm.get_weight_version() == "default" + # Step 1: Initialize weight transfer engine llm.init_weight_transfer_engine( WeightTransferInitRequest(init_info={"test_param": "flow_test"}) @@ -278,8 +281,15 @@ def test_full_weight_transfer_flow(): ) ) + assert llm.get_weight_version() == "default" + # Step 4: Finish weight update - llm.finish_weight_update() + llm.finish_weight_update("step-42") + + assert llm.get_weight_version() == "step-42" + + llm.update_weight_version("manual-version") + assert llm.get_weight_version() == "manual-version" # Verify the full flow completed def check_flow(self): diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 2e377e29253..adddf41ff4e 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -370,7 +370,7 @@ class VLLMWeightSyncClient(Protocol): def update_weights(self, update_info: dict[str, Any]) -> None: ... - def finish_weight_update(self) -> None: ... + def finish_weight_update(self, weight_version: str | None = None) -> None: ... class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 4f54a6e291e..12dd0c9eacc 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -77,8 +77,11 @@ class HTTPVLLMWeightSyncClient: "update_weights", {"update_info": _json_safe_update_info(update_info)} ) - def finish_weight_update(self) -> None: - self._post("finish_weight_update") + def finish_weight_update(self, weight_version: str | None = None) -> None: + json = ( + {"weight_version": weight_version} if weight_version is not None else None + ) + self._post("finish_weight_update", json) class RayVLLMWeightSyncClient: @@ -108,7 +111,11 @@ class RayVLLMWeightSyncClient: request = WeightTransferUpdateRequest(update_info=update_info) ray.get([h.update_weights.remote(request) for h in self.handles]) - def finish_weight_update(self) -> None: + def finish_weight_update(self, weight_version: str | None = None) -> None: import ray ray.get([h.finish_weight_update.remote() for h in self.handles]) + if weight_version is not None: + ray.get( + [h.update_weight_version.remote(weight_version) for h in self.handles] + ) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index ef3be178ac8..5a9b9f96d2c 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -267,6 +267,14 @@ class EngineClient(ABC): """Batched weight update for RL training.""" raise NotImplementedError - async def finish_weight_update(self) -> None: - """Finish the current weight update.""" + async def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + raise NotImplementedError + + async def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + raise NotImplementedError + + async def get_weight_version(self) -> str: + """Return the latest committed weight version.""" raise NotImplementedError diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index b3205728e49..4274819d988 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -885,9 +885,19 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): "update_weights", kwargs={"update_info": update_info_dict} ) - def finish_weight_update(self) -> None: - """Finish the current weight update.""" + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py index 8a2494a59df..392fcf56747 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -5,7 +5,7 @@ import json from http import HTTPStatus from typing import Annotated -from fastapi import APIRouter, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -203,11 +203,29 @@ async def update_weights(raw_request: Request): @router.post("/finish_weight_update") -async def finish_weight_update(raw_request: Request): - await engine_client(raw_request).finish_weight_update() +async def finish_weight_update( + raw_request: Request, + weight_version: Annotated[str | None, Body(embed=True)] = None, +): + await engine_client(raw_request).finish_weight_update(weight_version) return JSONResponse(content={"message": "Weight update finished"}) +@router.post("/update_weight_version") +async def update_weight_version( + raw_request: Request, + new_version: Annotated[str, Body(embed=True)], +): + await engine_client(raw_request).update_weight_version(new_version) + return JSONResponse(content={"success": True, "new_version": new_version}) + + +@router.get("/weight_info") +async def weight_info(raw_request: Request): + weight_version = await engine_client(raw_request).get_weight_version() + return JSONResponse(content={"weight_version": weight_version}) + + @router.get("/get_world_size") async def get_world_size( raw_request: Request, diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 922a8aa5982..5c2e01cf44b 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1106,6 +1106,16 @@ class AsyncLLM(EngineClient): "update_weights", kwargs={"update_info": request.update_info} ) - async def finish_weight_update(self) -> None: - """Finish the current weight update.""" + async def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" await self.collective_rpc("finish_weight_update") + if weight_version is not None: + await self.update_weight_version(weight_version) + + async def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + await self.engine_core.set_weight_version_async(new_version) + + async def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return await self.engine_core.get_weight_version_async() diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9817c474343..9917f810b5b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -125,6 +125,8 @@ class EngineCore: ) self.log_stats = log_stats + # Opaque weight version supplied by the caller. + self._weight_version = "default" # Setup Model. self.model_executor = executor_class(vllm_config) @@ -956,6 +958,13 @@ class EngineCore: ) -> list[_R]: return self.model_executor.collective_rpc(method, timeout, args, kwargs) + def set_weight_version(self, weight_version: str) -> None: + self._weight_version = weight_version + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self._weight_version + def preprocess_add_request(self, request: EngineCoreRequest) -> tuple[Request, int]: """Preprocess the request. diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 0aa4b6f3312..a6c232b2ab7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -176,9 +176,21 @@ class EngineCoreClient(ABC): def execute_dummy_batch(self) -> None: raise NotImplementedError + def set_weight_version(self, weight_version: str) -> None: + raise NotImplementedError + + def get_weight_version(self) -> str: + raise NotImplementedError + async def execute_dummy_batch_async(self) -> None: raise NotImplementedError + async def set_weight_version_async(self, weight_version: str) -> None: + raise NotImplementedError + + async def get_weight_version_async(self) -> str: + raise NotImplementedError + def abort_requests(self, request_ids: list[str]) -> None: raise NotImplementedError @@ -351,6 +363,12 @@ class InprocClient(EngineCoreClient): def execute_dummy_batch(self) -> None: self.engine_core.execute_dummy_batch() + def set_weight_version(self, weight_version: str) -> None: + self.engine_core.set_weight_version(weight_version) + + def get_weight_version(self) -> str: + return self.engine_core.get_weight_version() + def add_lora(self, lora_request: LoRARequest) -> bool: return self.engine_core.add_lora(lora_request) @@ -947,6 +965,12 @@ class SyncMPClient(MPClient): def execute_dummy_batch(self) -> None: self.call_utility("execute_dummy_batch") + def set_weight_version(self, weight_version: str) -> None: + self.call_utility("set_weight_version", weight_version) + + def get_weight_version(self) -> str: + return self.call_utility("get_weight_version") + def collective_rpc( self, method: str | Callable[..., _R], @@ -1199,6 +1223,12 @@ class AsyncMPClient(MPClient): async def execute_dummy_batch_async(self) -> None: await self.call_utility_async("execute_dummy_batch") + async def set_weight_version_async(self, weight_version: str) -> None: + await self.call_utility_async("set_weight_version", weight_version) + + async def get_weight_version_async(self) -> str: + return await self.call_utility_async("get_weight_version") + async def add_lora_async(self, lora_request: LoRARequest) -> bool: return await self.call_utility_async("add_lora", lora_request) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index ff86a1dffd9..17e40630859 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -425,6 +425,13 @@ class LLMEngine: ) -> list[_R]: return self.engine_core.collective_rpc(method, timeout, args, kwargs) + def set_weight_version(self, weight_version: str) -> None: + self.engine_core.set_weight_version(weight_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.engine_core.get_weight_version() + def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]: return self.collective_rpc("apply_model", args=(func,)) From 03a2d033673d8804b001bae41f30059273e1a669 Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 28 Jul 2026 14:38:20 +0800 Subject: [PATCH 10/67] [Bugfix] Respect cgroup memory limits on all platforms (#49966) Signed-off-by: chaunceyjiang Co-authored-by: Andreas Karatzas --- vllm/model_executor/model_loader/weight_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index db161a58988..e1897a77f10 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -694,12 +694,10 @@ def _get_checkpoints_size_bytes(files: list[str]) -> int: def _get_available_ram_bytes() -> int: - """Return available RAM, honoring cgroup limits on ROCm.""" + """Return available RAM, honoring cgroup limits.""" import psutil host_available = psutil.virtual_memory().available - if not current_platform.is_rocm(): - return host_available from vllm.utils.cpu_resource_utils import get_cgroup_memory_limit From b09688a6e77f0c061c08c67aeed56ad669c617dc Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 28 Jul 2026 15:16:21 +0800 Subject: [PATCH 11/67] [Bugfix][Spec Decode] Preserve draft buffers across level-2 sleep (#49774) Signed-off-by: aoshen02 Signed-off-by: vx120 <893600387@qq.com> Co-authored-by: aoshen02 Co-authored-by: vx120 <893600387@qq.com> Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Roger Wang --- vllm/v1/worker/gpu_worker.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 556b1e6c7d9..bbab98dbd7b 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -154,7 +154,7 @@ class Worker(WorkerBase): self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} - self._sleep_rebuild_draft_metadata_buffers = False + self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {} # Weight transfer engine is created in `load_model` once the model # is available, since the engine needs a reference to the model. @@ -200,10 +200,10 @@ class Worker(WorkerBase): name: buffer.cpu().clone() for name, buffer in model.named_buffers() } draft = self.get_draft_model() - inner = getattr(draft, "model", None) if draft is not None else None - self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr( - inner, "_build_fused_kv_buffers" - ) + if draft is not None: + self._sleep_saved_draft_buffers = { + name: buffer.cpu().clone() for name, buffer in draft.named_buffers() + } self._get_sleep_mode_backend().suspend(level) @@ -228,20 +228,21 @@ class Worker(WorkerBase): self._get_sleep_mode_backend().resume(tags) # Restore the buffers after level 2 sleep - if len(self._sleep_saved_buffers): + wake_weights = tags is None or "weights" in tags + if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): if name in self._sleep_saved_buffers: buffer.data.copy_(self._sleep_saved_buffers[name].data) self._sleep_saved_buffers = {} - if self._sleep_rebuild_draft_metadata_buffers: + if wake_weights and len(self._sleep_saved_draft_buffers): draft = self.get_draft_model() if draft is not None: - inner = getattr(draft, "model", None) - if inner is not None and hasattr(inner, "_build_fused_kv_buffers"): - inner._build_fused_kv_buffers() - self._sleep_rebuild_draft_metadata_buffers = False + for name, buffer in draft.named_buffers(): + if name in self._sleep_saved_draft_buffers: + buffer.data.copy_(self._sleep_saved_draft_buffers[name].data) + self._sleep_saved_draft_buffers = {} if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() From 99b57a4823d8fe22d3e249b40efa5d0f6fdfa2b1 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 28 Jul 2026 02:18:04 -0500 Subject: [PATCH 12/67] [CI][ROCm] Soft fail LoRA mirror (#50086) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/lora.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index a79196ffbd8..214d4a12bf7 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -16,6 +16,7 @@ steps: amd: dind: false device: mi300_1 + soft_fail: true working_dir: "/vllm-workspace/tests" timeout_in_minutes: 85 source_file_dependencies: From 61ac368021033dea54448d36ed98c7426e82cd18 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Tue, 28 Jul 2026 15:22:09 +0800 Subject: [PATCH 13/67] [Kimi-K3] Add AttnRes kernels (#50090) Signed-off-by: Thien Tran --- .buildkite/test_areas/models_basic.yaml | 12 + CMakeLists.txt | 22 + .../kimi_k3/attn_res_kernel.cu | 954 ++++++++++++++++++ csrc/libtorch_stable/ops.h | 11 + csrc/libtorch_stable/torch_bindings.cpp | 11 + tests/models/kimi_k3/test_amd_attn_res.py | 102 ++ tests/models/kimi_k3/test_attn_res.py | 193 ++++ vllm/_custom_ops.py | 27 + vllm/models/kimi_k3/amd/__init__.py | 2 + vllm/models/kimi_k3/amd/ops/__init__.py | 0 vllm/models/kimi_k3/amd/ops/attn_res.py | 132 +++ vllm/models/kimi_k3/nvidia/__init__.py | 2 + vllm/models/kimi_k3/nvidia/ops/__init__.py | 6 + vllm/models/kimi_k3/nvidia/ops/attn_res.py | 245 +++++ 14 files changed, 1719 insertions(+) create mode 100644 csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu create mode 100644 tests/models/kimi_k3/test_amd_attn_res.py create mode 100644 tests/models/kimi_k3/test_attn_res.py create mode 100644 vllm/models/kimi_k3/amd/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/attn_res.py create mode 100644 vllm/models/kimi_k3/nvidia/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/attn_res.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index af90308d6c3..a7c7aa9022d 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -61,6 +61,18 @@ steps: # FA4 kernel tests require SM100; the suite skips them elsewhere. - pytest -v -s models/inkling +- label: Kimi K3 Unit Tests (B200) + key: kimi-k3-unit-tests-b200 + timeout_in_minutes: 40 + device: b200-k8s + source_file_dependencies: + - vllm/models/kimi_k3/ + - csrc/libtorch_stable/kimi_k3/ + - tests/models/kimi_k3/ + commands: + # The native NVIDIA AttnRes kernel requires the SM100 family. + - pytest -v -s models/kimi_k3 + - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu depends_on: diff --git a/CMakeLists.txt b/CMakeLists.txt index 3a4e23ad0b3..cdda81ea46e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1079,6 +1079,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") set(MLA_ARCHS) endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS + "10.0f" "${CUDA_ARCHS}") + endif() + if(KIMI_K3_ATTN_RES_ARCHS) + set(KIMI_K3_ATTN_RES_SRC + "csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${KIMI_K3_ATTN_RES_SRC}" + CUDA_ARCHS "${KIMI_K3_ATTN_RES_ARCHS}") + set_property(SOURCE ${KIMI_K3_ATTN_RES_SRC} APPEND PROPERTY + COMPILE_OPTIONS + "$<$:--expt-relaxed-constexpr;--expt-extended-lambda;--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${KIMI_K3_ATTN_RES_SRC}") + message(STATUS + "Building Kimi K3 AttnRes for archs: ${KIMI_K3_ATTN_RES_ARCHS}") + endif() + # Hadacore kernels cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}") if(HADACORE_ARCHS) @@ -1120,6 +1138,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_COOPERATIVE_TOPK=1) endif() + if(KIMI_K3_ATTN_RES_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_KIMI_K3_ATTN_RES=1) + endif() # Needed by CUTLASS kernels target_compile_definitions(_C_stable_libtorch PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu new file mode 100644 index 00000000000..eb4dcf6bf5a --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu @@ -0,0 +1,954 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +// Production AttnRes forward for Blackwell (SM100). +// +// Warp-specialized online softmax + residual + RMSNorm: +// - 1 producer warp issues cp.async.bulk row loads into shared memory. +// - 8 consumer warps compute reductions and output. +// - Q=res_weight*rms_weight remains in registers across persistent tokens. +// - V rows are converted once and cached as FP32 in TMEM between passes. +// +// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major +// block residual storage. + +#include "../torch_utils.h" + +#include +#include +#include +#include +#include + +using bf16_t = __nv_bfloat16; + +namespace sm100 { +namespace fwd_prod_v2 { + +constexpr int K_TILE = 1024; +constexpr int N_CHUNK_DEFAULT = 4; +constexpr int CHUNK_DEPTH = 2; +constexpr int BLK = 288; // 1 producer warp + 8 consumer warps +constexpr int CONSUMER_THREADS = BLK - 32; // 256 +constexpr int CONSUMER_WARPS = CONSUMER_THREADS / 32; +constexpr int CONSUMER_GROUPS = 2; // two 128-thread consumer groups +constexpr int CONSUMER_THREADS_PER_GROUP = CONSUMER_THREADS / CONSUMER_GROUPS; +constexpr int FIRST_USER_NAMED_BARRIER = 8; + +__device__ __forceinline__ const bf16_t* residual_addr( + const bf16_t* block_res, const bf16_t* layer_res, int source, int N, + int token, int block_stride_m, int block_stride_r, int H) { + if (source < N - 1) { + return block_res + static_cast(token) * block_stride_m + + source * block_stride_r; + } + return layer_res + static_cast(token) * H; +} + +__device__ __forceinline__ uint32_t elect_one_sync() { + uint32_t pred = 0; + uint32_t laneid = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %2;\n" + "@%%px mov.s32 %1, 1;\n" + " mov.s32 %0, %%rx;\n" + "}\n" + : "+r"(laneid), "+r"(pred) + : "r"(0xffffffff)); + return pred; +} + +__device__ __forceinline__ void mbarrier_init(uint64_t& barrier, + int thread_count) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(barrier_addr), + "r"(thread_count)); +} + +__device__ __forceinline__ void mbarrier_expect_tx(uint64_t& barrier, + uint32_t bytes) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"( + barrier_addr), + "r"(bytes)); +} + +__device__ __forceinline__ void mbarrier_wait(uint64_t& barrier, int phase) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .pred p;\n" + "WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n" + "@p bra DONE;\n" + "bra WAIT;\n" + "DONE:\n" + "}\n" ::"r"(barrier_addr), + "r"(phase)); +} + +__device__ __forceinline__ void mbarrier_arrive(uint64_t& barrier) { + uint32_t const barrier_addr = + static_cast(__cvta_generic_to_shared(&barrier)); + asm volatile( + "{\n" + ".reg .b64 state;\n" + "mbarrier.arrive.shared::cta.b64 state, [%0];\n" + "}\n" ::"r"(barrier_addr)); +} + +__device__ __forceinline__ void fence_mbarrier_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +__device__ __forceinline__ void named_barrier_sync(uint32_t num_threads, + uint32_t user_barrier_id) { + asm volatile( + "bar.sync %0, %1;" ::"r"(user_barrier_id + FIRST_USER_NAMED_BARRIER), + "r"(num_threads) + : "memory"); +} + +__device__ __forceinline__ void tmem_allocate(int num_columns, uint32_t* dst) { + uint32_t const dst_addr = + static_cast(__cvta_generic_to_shared(dst)); + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"( + dst_addr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_free(uint32_t tmem_ptr, int num_columns) { + asm volatile( + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(tmem_ptr), + "r"(num_columns)); +} + +__device__ __forceinline__ void tmem_release_allocation_lock() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); +} + +__device__ __forceinline__ void tmem_store_wait() { + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); +} + +template +__device__ __forceinline__ void tmem_load(uint32_t src_addr, T* dst) { + uint32_t* values = reinterpret_cast(dst); + if constexpr (N == 8) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x8.b32" + "{%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]), + "=r"(values[4]), "=r"(values[5]), "=r"(values[6]), "=r"(values[7]) + : "r"(src_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x4.b32" + "{%0, %1, %2, %3}, [%4];\n" + : "=r"(values[0]), "=r"(values[1]), "=r"(values[2]), "=r"(values[3]) + : "r"(src_addr)); + } +} + +template +__device__ __forceinline__ void tmem_store(uint32_t dst_addr, T* src) { + uint32_t* values = reinterpret_cast(src); + if constexpr (N == 8) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + "[%8], {%0, %1, %2, %3, %4, %5, %6, %7};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(values[4]), + "r"(values[5]), "r"(values[6]), "r"(values[7]), "r"(dst_addr)); + } else { + static_assert(N == 4, "AttnRes TMEM helpers support x4 and x8"); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x4.b32" + "[%4], {%0, %1, %2, %3};\n" ::"r"(values[0]), + "r"(values[1]), "r"(values[2]), "r"(values[3]), "r"(dst_addr)); + } +} + +__device__ __forceinline__ float2 float2_add(const float2& a, const float2& b) { + float2 result; + asm volatile("add.rn.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b))); + return result; +} + +__device__ __forceinline__ float2 float2_mul(const float2& a, const float2& b) { + float2 result; + asm volatile("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b))); + return result; +} + +__device__ __forceinline__ float2 float2_fma(const float2& a, const float2& b, + const float2& c) { + float2 result; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast(result)) + : "l"(reinterpret_cast(a)), + "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); + return result; +} + +template +struct FwdSmemPlan { + alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; + alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; + alignas(16) uint64_t bar_output_norm_ready; + alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; + uint32_t tmem_base; +}; + +__device__ __forceinline__ void cp_async_bulk(void* smem_dst, + const void* gmem_src, int bytes, + uint64_t& mbar) { + uint32_t const s = static_cast(__cvta_generic_to_shared(smem_dst)); + uint32_t const m = static_cast(__cvta_generic_to_shared(&mbar)); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], " + "[%1], %2, [%3];\n" ::"r"(s), + "l"(gmem_src), "r"(bytes), "r"(m) + : "memory"); +} + +template +__global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( + const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res, + const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w, + const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int N, int T, + int B, int block_stride_m, int block_stride_r, float rms_eps, + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 + constexpr float LOG2_E = 1.4426950408889634f; + constexpr int N_CHUNK = NC; + // The two-source specialization only consumes half of the TMEM columns. + constexpr int TMEM_COLS_ALLOC = NC == 2 ? 128 : 256; + constexpr int NUM_BUFS = CHUNK_DEPTH * NC; + constexpr int NHT = H / K_TILE; + constexpr int SLICES_PER_GROUP = + (NHT + CONSUMER_GROUPS - 1) / CONSUMER_GROUPS; + constexpr int VEC = 8; + constexpr int ACC_PER_THREAD = H == 7168 ? 28 : SLICES_PER_GROUP * VEC; + constexpr int TMEM_V_COLS_PER_GROUP = SLICES_PER_GROUP * N_CHUNK * VEC; + constexpr int TMEM_V_COLS_TOTAL = CONSUMER_GROUPS * TMEM_V_COLS_PER_GROUP; + static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC); + static_assert(H >= 4096 && H <= 8192); + static_assert(H % K_TILE == 0); + + const int tid = threadIdx.x; + const int wid = tid >> 5; + const int lane = tid & 31; + const int TB = T * B; + const int num_ctas = gridDim.x; + const int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; + + const int comp_wid = wid - 1; + const int comp_tid = tid - 32; + const int group = (comp_wid >= 4) ? 1 : 0; + const int ct_in_group = + (comp_tid >= 0) ? (comp_tid & (CONSUMER_THREADS_PER_GROUP - 1)) : -1; + const int k_local = ct_in_group * VEC; + + constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); + constexpr size_t DELTA_BYTES = + HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0; + constexpr size_t OUTPUT_NORM_BYTES = + OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0; + extern __shared__ __align__(16) char smem_raw[]; + bf16_t* v_bufs = reinterpret_cast(smem_raw); // [NUM_BUFS][H] + bf16_t* delta_bufs = reinterpret_cast(smem_raw + V_BYTES); + bf16_t* output_norm_buf = + reinterpret_cast(smem_raw + V_BYTES + DELTA_BYTES); + FwdSmemPlan& plan = *reinterpret_cast*>( + smem_raw + V_BYTES + DELTA_BYTES + OUTPUT_NORM_BYTES); + + auto slot_of = [](long long gci, int n) { + return (int)(gci % CHUNK_DEPTH) * N_CHUNK + n; + }; + auto phase_of = [](long long gci) { return (int)((gci / CHUNK_DEPTH) & 1); }; + auto buf_ptr = [&](int slot) -> bf16_t* { return v_bufs + slot * H; }; + auto delta_buf_ptr = [&](int chunk_slot) -> bf16_t* { + return delta_bufs + chunk_slot * H; + }; + + if (wid == 0 && elect_one_sync()) { + #pragma unroll + for (int i = 0; i < CHUNK_DEPTH; i++) { + mbarrier_init(plan.bar_ready[i], 1); + mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS); + } + if constexpr (OUTPUT_NORM_IN_SMEM) { + mbarrier_init(plan.bar_output_norm_ready, 1); + } + fence_mbarrier_init(); + } + + // gdc wait BEFORE tmem alloc + cudaGridDependencySynchronize(); + + if (wid == 1) { + tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base); + if constexpr (RELEASE_TMEM) { + tmem_release_allocation_lock(); + } + } + __syncthreads(); + + if constexpr (OUTPUT_NORM_IN_SMEM) { + if (wid == 0 && elect_one_sync()) { + mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t)); + cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t), + plan.bar_output_norm_ready); + } + } + + const uint32_t my_v_tmem = + comp_tid >= 0 ? plan.tmem_base + group * TMEM_V_COLS_PER_GROUP : 0; + float q_cache[ACC_PER_THREAD]; + if (comp_tid >= 0) { + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + #pragma unroll + for (int j = 0; j < 4; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + #pragma unroll + for (int j = 0; j < VEC; j++) { + int h = h_base + j; + q_cache[si * VEC + j] = + __bfloat162float(rms_w[h]) * __bfloat162float(res_w[h]); + } + } + } + + if (wid == 0) { + if (elect_one_sync()) { + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + const int t = tb / B; + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pc = phase_of(gci); + mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1); + int transaction_bytes = an * H * (int)sizeof(bf16_t); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + transaction_bytes += H * (int)sizeof(bf16_t); + } + } + mbarrier_expect_tx(plan.bar_ready[chunk_slot], transaction_bytes); + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n >= an) continue; + int slot = slot_of(gci, n); + const bf16_t* src = + residual_addr(block_res, layer_res, ns + n, N, t, + block_stride_m, block_stride_r, H); + cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (prefix_n >= 0 && prefix_n < an) { + cp_async_bulk(delta_buf_ptr(chunk_slot), + delta + (long long)tb * H, H * sizeof(bf16_t), + plan.bar_ready[chunk_slot]); + } + } + } + } + } + } else { + float acc32[ACC_PER_THREAD] = {}; + float eps_cache; + asm volatile("mov.b32 %0, %1;" : "=f"(eps_cache) : "f"(rms_eps)); + + long long gci = 0; + for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { + float m_running = -FLT_MAX; + float s_running = 0.f; + #pragma unroll + for (int i = 0; i < ACC_PER_THREAD; i++) { + acc32[i] = 0.f; + } + + for (int ci = 0; ci < num_chunks; ci++, gci++) { + int ns = ci * N_CHUNK; + int an = min(N_CHUNK, N - ns); + int chunk_slot = (int)(gci % CHUNK_DEPTH); + int pr = phase_of(gci); + mbarrier_wait(plan.bar_ready[chunk_slot], pr); + + float2 sq_local[N_CHUNK] = {}; + float2 dot_local[N_CHUNK] = {}; + + auto pass_A_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = + 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + const float* qv = &q_cache[si * VEC]; + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int2 vp = + *reinterpret_cast(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = + delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < 2; j++) { + auto delta2 = *reinterpret_cast( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[2] = {__bfloat1622float2(v2[0]), + __bfloat1622float2(v2[1])}; + tmem_store<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + sq_local[n] = float2_fma(f[0], f[0], sq_local[n]); + sq_local[n] = float2_fma(f[1], f[1], sq_local[n]); + dot_local[n] = + float2_fma(f[0], make_float2(qv[0], qv[1]), dot_local[n]); + dot_local[n] = + float2_fma(f[1], make_float2(qv[2], qv[3]), dot_local[n]); + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + const float* qv = &q_cache[si * VEC]; + + #pragma unroll + for (int n = 0; n < AN; n++) { + int slot = slot_of(gci, n); + int4 vp = *reinterpret_cast(buf_ptr(slot) + h_base); + auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); + if constexpr (HAS_DELTA) { + int prefix_n = N - 1 - ns; + if (n == prefix_n) { + const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + auto delta2 = *reinterpret_cast( + delta_ptr + 2 * j); + v2[j] = __hadd2(v2[j], delta2); + } + *reinterpret_cast(layer_res + (long long)tb * H + + h_base) = vp; + } + } + float2 f[4] = { + __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]), + __bfloat1622float2(v2[2]), __bfloat1622float2(v2[3])}; + tmem_store(my_v_tmem + (si * N_CHUNK + n) * VEC, f); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + sq_local[n] = float2_fma(f[j], f[j], sq_local[n]); + dot_local[n] = float2_fma( + f[j], make_float2(qv[2 * j], qv[2 * j + 1]), dot_local[n]); + } + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_A_body(std::integral_constant{}); + break; + case 3: + pass_A_body(std::integral_constant{}); + break; + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_A_body(std::integral_constant{}); + break; + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_A_body(std::integral_constant{}); + break; + case 1: + pass_A_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } + if (lane == 0) { + mbarrier_arrive(plan.bar_consumed[chunk_slot]); + } + tmem_store_wait(); + + float2 reduce_pair[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + reduce_pair[n] = make_float2(sq_local[n].x + sq_local[n].y, + dot_local[n].x + dot_local[n].y); + } + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + uint64_t packed = reinterpret_cast(reduce_pair[n]); + packed = __shfl_xor_sync(0xffffffff, packed, offset); + float2 other = reinterpret_cast(packed); + reduce_pair[n] = float2_add(reduce_pair[n], other); + } + } + if (lane == 0) { + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + plan.ws_stats[comp_wid][n] = reduce_pair[n]; + } + } + named_barrier_sync(CONSUMER_THREADS, 0); + + float local_rsig = 0.f; + float local_logit = 0.f; + int stat_n = lane / CONSUMER_WARPS; + int stat_w = lane % CONSUMER_WARPS; + float2 totals = {}; + if (stat_n < N_CHUNK) { + totals = plan.ws_stats[stat_w][stat_n]; + } + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + totals.x += + __shfl_down_sync(0xffffffff, totals.x, offset, CONSUMER_WARPS); + totals.y += + __shfl_down_sync(0xffffffff, totals.y, offset, CONSUMER_WARPS); + } + if (stat_n < N_CHUNK && stat_w == 0) { + local_rsig = rsqrtf(totals.x / H + eps_cache); + local_logit = totals.y * local_rsig; + } + float logit_n[N_CHUNK]; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + logit_n[n] = __shfl_sync(0xffffffff, local_logit, n * CONSUMER_WARPS); + } + + float m_chunk = -FLT_MAX; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) m_chunk = fmaxf(m_chunk, logit_n[n]); + } + float m_new = fmaxf(m_running, m_chunk); + float corr = exp2f((m_running - m_new) * LOG2_E); + float w_n[N_CHUNK] = {}; + float w_sum = 0.f; + #pragma unroll + for (int n = 0; n < N_CHUNK; n++) { + if (n < an) { + w_n[n] = exp2f((logit_n[n] - m_new) * LOG2_E); + w_sum += w_n[n]; + } + } + + auto pass_B_body = [&](auto AN_TOK) { + constexpr int AN = decltype(AN_TOK)::value; + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + float2 corr2 = make_float2(corr, corr); + float2 a[2]; + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load<4>(my_v_tmem + (si * N_CHUNK + n) * VEC, + f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + float2 corr2 = make_float2(corr, corr); + float2 a[VEC / 2]; + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + a[j] = float2_mul(old, corr2); + } + float2 f_cache[AN][VEC / 2]; + #pragma unroll + for (int n = 0; n < AN; n++) { + tmem_load(my_v_tmem + (si * N_CHUNK + n) * VEC, f_cache[n]); + } + #pragma unroll + for (int n = 0; n < AN; n++) { + float2 wn = make_float2(w_n[n], w_n[n]); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + a[j] = float2_fma(wn, f_cache[n][j], a[j]); + } + } + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + acc32[si * VEC + 2 * j] = a[j].x; + acc32[si * VEC + 2 * j + 1] = a[j].y; + } + } + }; + if constexpr (NC == 4) { + switch (an) { + case 4: + pass_B_body(std::integral_constant{}); + break; + case 3: + pass_B_body(std::integral_constant{}); + break; + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else if constexpr (NC == 3) { + switch (an) { + case 3: + pass_B_body(std::integral_constant{}); + break; + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } else { + static_assert(NC == 2); + switch (an) { + case 2: + pass_B_body(std::integral_constant{}); + break; + case 1: + pass_B_body(std::integral_constant{}); + break; + default: + __builtin_unreachable(); + } + } + + s_running = s_running * corr + w_sum; + m_running = m_new; + } + + float inv_s = 1.f / s_running; + bf16_t* out_ptr = output + (long long)tb * H; + float2 output_sq_pair = {}; + // When output RMSNorm is fused, the softmax denominator cancels: + // (acc / s) * rsqrt(mean((acc / s)^2) + eps) + // = acc * rsqrt(mean(acc^2) + eps * s^2). + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 old = make_float2(acc32[si * VEC + 2 * j], + acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast(out_ptr + h_base) = packed; + } + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* ov2 = reinterpret_cast<__nv_bfloat162*>(&packed); + float2 inv2 = make_float2(inv_s, inv_s); + #pragma unroll + for (int j = 0; j < VEC / 2; j++) { + float2 old = + make_float2(acc32[si * VEC + 2 * j], acc32[si * VEC + 2 * j + 1]); + if constexpr (HAS_OUTPUT_NORM) { + output_sq_pair = float2_fma(old, old, output_sq_pair); + } else { + float2 mixed = float2_mul(old, inv2); + ov2[j] = __float22bfloat162_rn(mixed); + } + } + if constexpr (!HAS_OUTPUT_NORM) { + *reinterpret_cast(out_ptr + h_base) = packed; + } + } + + if constexpr (HAS_OUTPUT_NORM) { + if constexpr (OUTPUT_NORM_IN_SMEM) { + // The immutable weight copy is acquired once, at its first use. + if (tb == blockIdx.x) { + mbarrier_wait(plan.bar_output_norm_ready, 0); + } + } + float output_sq = output_sq_pair.x + output_sq_pair.y; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + output_sq += __shfl_xor_sync(0xffffffff, output_sq, offset); + } + if (lane == 0) { + plan.ws_stats[comp_wid][0] = make_float2(output_sq, 0.f); + } + named_barrier_sync(CONSUMER_THREADS, 0); + float total_sq = lane < CONSUMER_WARPS ? plan.ws_stats[lane][0].x : 0.f; + #pragma unroll + for (int offset = CONSUMER_WARPS / 2; offset > 0; offset >>= 1) { + total_sq += + __shfl_down_sync(0xffffffff, total_sq, offset, CONSUMER_WARPS); + } + if (lane == 0) { + total_sq = + rsqrtf(total_sq / H + output_norm_eps * s_running * s_running); + } + float output_rsigma = __shfl_sync(0xffffffff, total_sq, 0); + #pragma unroll + for (int si = 0; si < SLICES_PER_GROUP; si++) { + if constexpr (H == 7168) { + if (si == SLICES_PER_GROUP - 1) { + int h_base = 6 * K_TILE + group * (K_TILE / 2) + ct_in_group * 4; + uint2 packed; + auto* values = reinterpret_cast(&packed); + #pragma unroll + for (int j = 0; j < 4; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = __float2bfloat16(acc32[si * VEC + j] * + output_rsigma * weight); + } + *reinterpret_cast(out_ptr + h_base) = packed; + continue; + } + } + int dt = si * CONSUMER_GROUPS + group; + if (dt >= NHT) continue; + int h_base = dt * K_TILE + k_local; + uint4 packed; + auto* values = reinterpret_cast(&packed); + #pragma unroll + for (int j = 0; j < VEC; j++) { + const bf16_t* weight_ptr = + OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; + float weight = __bfloat162float(weight_ptr[h_base + j]); + values[j] = + __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); + } + *reinterpret_cast(out_ptr + h_base) = packed; + } + } + } + } + + cudaTriggerProgrammaticLaunchCompletion(); + __syncthreads(); + if (wid == 1) { + tmem_free(plan.tmem_base, TMEM_COLS_ALLOC); + } +#else + if (threadIdx.x == 0) { + printf("attn_res_fwd_online_v2_kernel requires sm_10x\n"); + } +#endif +} + +template +static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual, + const bf16_t* delta, const bf16_t* res_weight, + const bf16_t* rms_weight, bf16_t* output, int N, int T, + int B, float rms_eps, int num_sm, cudaStream_t stream, + const bf16_t* output_norm_weight = nullptr, + float output_norm_eps = 0.f, int block_stride_m = 0, + int block_stride_r = 0) { + constexpr size_t smem_size = + ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) + + (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) + + sizeof(FwdSmemPlan) + 15) & + ~size_t(15); + auto kernel = + &attn_res_fwd_online_v2_kernel; + static bool attrs_set = false; + if (!attrs_set) { + if (smem_size > 48 * 1024) { + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + } + attrs_set = true; + } + int grid = RELEASE_TMEM ? num_sm * 2 : num_sm; + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = BLK; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta, + res_weight, rms_weight, output, N, T, B, block_stride_m, + block_stride_r, rms_eps, output_norm_weight, + output_norm_eps); +} + +} // namespace fwd_prod_v2 +} // namespace sm100 + +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps) { + int const num_tokens = static_cast(prefix.size(0)); + int const device = prefix.get_device_index(); + torch::stable::accelerator::DeviceGuard const device_guard(device); + cudaDeviceProp const* properties = get_device_prop(); + STD_TORCH_CHECK(properties->major == 10, + "Kimi K3 AttnRes requires the SM100 family"); + + using namespace sm100::fwd_prod_v2; + // Two-source chunks and two resident CTAs are beneficial once setup is + // amortized by the long, full eight-block prefill workload. + if (num_blocks == 8 && num_tokens >= 4096) { + launch_fwd<7168, 2, true, true, true, true>( + static_cast(blocks.data_ptr()), + static_cast(prefix.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(qk_weight.data_ptr()), + static_cast(norm_weight.data_ptr()), + static_cast(output.data_ptr()), + static_cast(num_blocks) + 1, num_tokens, 1, + static_cast(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast(output_norm_weight.data_ptr()), + static_cast(output_norm_eps), static_cast(blocks.stride(0)), + static_cast(blocks.stride(1))); + } else { + launch_fwd<7168, 4, false, true, true, true>( + static_cast(blocks.data_ptr()), + static_cast(prefix.data_ptr()), + static_cast(delta.data_ptr()), + static_cast(qk_weight.data_ptr()), + static_cast(norm_weight.data_ptr()), + static_cast(output.data_ptr()), + static_cast(num_blocks) + 1, num_tokens, 1, + static_cast(eps), properties->multiProcessorCount, + get_current_cuda_stream(device), + static_cast(output_norm_weight.data_ptr()), + static_cast(output_norm_eps), static_cast(blocks.stride(0)), + static_cast(blocks.stride(1))); + } + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 AttnRes kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 3834bea5857..5a9c91d4563 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -315,6 +315,17 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional index_q_out, const std::string& kv_cache_dtype, bool skip_index_branch); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES +void kimi_k3_attn_res(torch::stable::Tensor& prefix, + torch::stable::Tensor const& delta, + torch::stable::Tensor const& blocks, + torch::stable::Tensor const& norm_weight, + torch::stable::Tensor const& qk_weight, + torch::stable::Tensor const& output_norm_weight, + torch::stable::Tensor& output, int64_t num_blocks, + double eps, double output_norm_eps); +#endif + // Sampler kernels (shared CUDA/ROCm) void apply_repetition_penalties_( torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 0a475d02c6f..c364e211474 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -468,6 +468,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int block_size, Tensor!? q_out, Tensor!? index_q_out, " "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.def( + "kimi_k3_attn_res(" + "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " + "int num_blocks, float eps, float output_norm_eps) -> ()"); +#endif + // Apply repetition penalties to logits in-place. ops.def( "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " @@ -693,6 +701,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES + ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); +#endif // Sampler kernels (shared CUDA/ROCm) ops.impl("apply_repetition_penalties_", diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py new file mode 100644 index 00000000000..f6dfaa422b3 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AMD AttnRes requires ROCm", +) + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + hidden_size = prefix.shape[-1] + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (hidden_size,), norm_weight, eps) + probs = (keys @ qk_weight).softmax(dim=-1) + return torch.matmul(probs.unsqueeze(1), values).squeeze(1) + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "block_capacity", + "hidden_size", + "row_padding", + ), + [ + pytest.param(0, 3, 5, 128, 0, id="empty"), + pytest.param(1, 1, 2, 128, 0, id="decode-single"), + pytest.param(17, 4, 6, 1024, 7, id="decode-padded"), + pytest.param(320, 8, 10, 7168, 0, id="prefill-full"), + ], +) +def test_amd_attn_res_matches_reference( + num_tokens: int, + num_blocks: int, + block_capacity: int, + hidden_size: int, + row_padding: int, +) -> None: + eps = 1e-5 + prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding) + blocks = _randn_with_row_padding( + num_tokens, + block_capacity, + hidden_size, + padding=row_padding, + ) + norm_weight = 1 + 0.1 * torch.randn( + hidden_size, device="cuda", dtype=torch.bfloat16 + ) + qk_weight = ( + torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5 + ) + expected = _reference( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + original_prefix = prefix.clone() + original_blocks = blocks.clone() + + actual = attn_res( + prefix, + blocks, + norm_weight, + qk_weight, + num_blocks, + eps, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.shape == prefix.shape + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py new file mode 100644 index 00000000000..69c9213647a --- /dev/null +++ b/tests/models/kimi_k3/test_attn_res.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.platforms import current_platform + +HIDDEN_SIZE = 7168 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: + storage = torch.randn( + *shape[:-1], + shape[-1] + padding, + device="cuda", + dtype=torch.bfloat16, + ) + return storage[..., : shape[-1]] + + +def _reference( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if delta is not None: + prefix = prefix + delta + values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1) + keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS) + probs = (keys @ qk_weight).softmax(dim=-1) + output = torch.matmul(probs.unsqueeze(1), values).squeeze(1) + if output_norm_weight is not None: + output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS) + return output, prefix + + +@pytest.mark.parametrize( + ( + "num_tokens", + "num_blocks", + "row_padding", + "write_block", + "has_delta", + "backend", + ), + [ + pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"), + pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"), + pytest.param(17, 5, 7, True, False, "triton", id="triton-write"), + pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), + pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), + pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), + pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), + pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), + pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), + ], +) +def test_attn_res( + num_tokens: int, + num_blocks: int, + row_padding: int, + write_block: bool, + has_delta: bool, + backend: str, +): + if backend == "nvidia" and not current_platform.is_device_capability_family(100): + pytest.skip("NVIDIA AttnRes requires the SM100 family") + + prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + delta = ( + _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding) + if has_delta + else None + ) + blocks = _randn_with_row_padding( + num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding + ) + norm_weight = 1 + 0.1 * torch.randn( + 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 = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + original_blocks = blocks.clone() + expected, expected_prefix = _reference( + prefix.clone(), + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + block_write_idx = num_blocks if write_block else -1 + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + block_write_idx, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) + if write_block: + original_blocks[:, block_write_idx].copy_(expected_prefix) + torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) + assert actual.is_contiguous() + + +@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1)) +def test_attn_res_block_counts(num_blocks: int): + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + blocks = torch.randn( + 1, 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, _ = _reference( + prefix.clone(), + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + ) + + actual = attn_res( + prefix, + None, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + -1, + EPS, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def test_attn_res_without_output_norm(): + prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.randn_like(prefix) + blocks = torch.randn( + 7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + expected, _ = _reference( + prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS + ) + + actual = attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + None, + MAX_BLOCKS, + -1, + EPS, + 0.0, + ) + + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 07a3a583f95..b1d2f1344d6 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2746,6 +2746,33 @@ def concat_and_cache_mla( ) +def kimi_k3_attn_res( + prefix: torch.Tensor, + delta: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor, + num_blocks: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + output = torch.empty_like(prefix) + torch.ops._C.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + num_blocks, + eps, + output_norm_eps, + ) + return output + + def concat_and_cache_mla_rope_fused( positions: torch.Tensor, q_pe: torch.Tensor, diff --git a/vllm/models/kimi_k3/amd/__init__.py b/vllm/models/kimi_k3/amd/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/amd/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/__init__.py b/vllm/models/kimi_k3/amd/ops/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/models/kimi_k3/amd/ops/attn_res.py b/vllm/models/kimi_k3/amd/ops/attn_res.py new file mode 100644 index 00000000000..c00a3422ca1 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/attn_res.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + eps: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + input_qk_weight = tl.load(norm_weight_ptr + d_offsets, mask=d_mask, other=0.0).to( + tl.float32 + ) * tl.load(qk_weight_ptr + d_offsets, mask=d_mask, other=0.0).to(tl.float32) + + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + num_sources = num_blocks + 1 + + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + block_values = tl.load( + block_ptrs, + mask=(source_mask[:, None] & ~is_prefix[:, None] & d_mask[None, :]), + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + values = tl.where(is_prefix[:, None], prefix[None, :], block_values) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + output = mixed / denominator + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + num_blocks: int, + eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert 0 < num_blocks <= blocks.shape[1] + assert blocks.shape[0] == num_tokens + assert norm_weight.numel() == hidden_size + assert qk_weight.numel() == hidden_size + assert prefix.stride(-1) == 1 + assert blocks.stride(-1) == 1 + assert norm_weight.stride(-1) == 1 + assert qk_weight.stride(-1) == 1 + + output = prefix.new_empty(prefix.shape) + if num_tokens == 0: + return output + + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + blocks, + norm_weight, + qk_weight, + output, + prefix.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + eps, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/__init__.py b/vllm/models/kimi_k3/nvidia/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/__init__.py b/vllm/models/kimi_k3/nvidia/ops/__init__.py new file mode 100644 index 00000000000..bbaee887ba2 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .attn_res import attn_res + +__all__ = ["attn_res"] diff --git a/vllm/models/kimi_k3/nvidia/ops/attn_res.py b/vllm/models/kimi_k3/nvidia/ops/attn_res.py new file mode 100644 index 00000000000..01078d6c01d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/attn_res.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + + +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_attn_res_triton_warmup_profiles( + max_blocks: int, +) -> tuple[tuple[int, bool, int, bool], ...]: + """Return the small-batch profiles that bypass the native kernel.""" + profiles = [ + (num_blocks, False, -1, True) for num_blocks in range(2, max_blocks + 1) + ] + profiles.extend( + (block_write_idx, True, block_write_idx, True) + for block_write_idx in range(2, max_blocks) + ) + profiles.append((max_blocks, True, -1, False)) + return tuple(profiles) + + +@triton.jit +def _attn_res_kernel( + prefix_ptr, + delta_ptr, + blocks_ptr, + norm_weight_ptr, + qk_weight_ptr, + output_norm_weight_ptr, + output_ptr, + stride_prefix_m: tl.constexpr, + stride_delta_m: tl.constexpr, + stride_block_m: tl.constexpr, + stride_block_r: tl.constexpr, + stride_output_m: tl.constexpr, + num_blocks: tl.constexpr, + hidden_size: tl.constexpr, + block_write_idx: tl.constexpr, + eps: tl.constexpr, + output_norm_eps: tl.constexpr, + HAS_DELTA: tl.constexpr, + WRITE_BLOCK: tl.constexpr, + APPLY_OUTPUT_NORM: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_D: tl.constexpr, + launch_pdl: tl.constexpr, +): + row_idx = tl.program_id(0).to(tl.int64) + d_offsets = tl.max_contiguous(tl.arange(0, BLOCK_D), BLOCK_D) + d_mask = d_offsets < hidden_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + updated_prefix = tl.load( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + if HAS_DELTA: + delta = tl.load( + delta_ptr + row_idx * stride_delta_m + d_offsets, + mask=d_mask, + other=0.0, + ).to(tl.float32) + updated_prefix += delta + # Match the BF16 prefix-add result before using it as a residual source. + updated_prefix = updated_prefix.to(prefix_ptr.dtype.element_ty).to(tl.float32) + tl.store( + prefix_ptr + row_idx * stride_prefix_m + d_offsets, + updated_prefix, + mask=d_mask, + ) + if WRITE_BLOCK: + tl.store( + blocks_ptr + + row_idx * stride_block_m + + block_write_idx * stride_block_r + + d_offsets, + updated_prefix, + mask=d_mask, + ) + # With only the prefix source, the AttnRes softmax is exactly one. + if num_blocks == 0: + mixed = updated_prefix + else: + # Reloading avoids keeping the full prefix vector live across the loop. + if HAS_DELTA: + tl.debug_barrier() + input_qk_weight = tl.load( + norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) * tl.load( + qk_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + max_logit = tl.full((), -float("inf"), tl.float32) + denominator = tl.zeros((), tl.float32) + mixed = tl.zeros((BLOCK_D,), tl.float32) + + num_sources = num_blocks + 1 + for source_tile in range(tl.cdiv(num_sources, BLOCK_L)): + source_offsets = source_tile * BLOCK_L + tl.arange(0, BLOCK_L) + source_mask = source_offsets < num_sources + is_prefix = source_offsets == num_blocks + block_ptrs = ( + blocks_ptr + + row_idx * stride_block_m + + source_offsets[:, None] * stride_block_r + + d_offsets[None, :] + ) + prefix_ptrs = ( + prefix_ptr + + row_idx * stride_prefix_m + + source_offsets[:, None] * 0 + + d_offsets[None, :] + ) + value_ptrs = tl.where(is_prefix[:, None], prefix_ptrs, block_ptrs) + values = tl.load( + value_ptrs, + mask=source_mask[:, None] & d_mask[None, :], + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + reciprocal_std = tl.rsqrt( + tl.sum(values * values, axis=1) * (1.0 / hidden_size) + eps + ) + logits = tl.sum(values * input_qk_weight[None, :], axis=1) * reciprocal_std + scores = tl.where(source_mask, logits, -float("inf")) + + new_max_logit = tl.maximum(max_logit, tl.max(scores, axis=0)) + old_scale = tl.exp(max_logit - new_max_logit) + block_scales = tl.exp(scores - new_max_logit) + denominator = denominator * old_scale + tl.sum(block_scales, axis=0) + mixed = mixed * old_scale + tl.sum(block_scales[:, None] * values, axis=0) + max_logit = new_max_logit + + mixed /= denominator + output = mixed + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + if APPLY_OUTPUT_NORM: + output_reciprocal_std = tl.rsqrt( + tl.sum(tl.where(d_mask, mixed * mixed, 0.0), axis=0) * (1.0 / hidden_size) + + output_norm_eps + ) + output_norm_weight = tl.load( + output_norm_weight_ptr + d_offsets, mask=d_mask, other=0.0 + ).to(tl.float32) + output = mixed * output_reciprocal_std * output_norm_weight + tl.store( + output_ptr + row_idx * stride_output_m + d_offsets, + output, + mask=d_mask, + ) + + +def attn_res( + prefix: torch.Tensor, + delta: torch.Tensor | None, + blocks: torch.Tensor, + norm_weight: torch.Tensor, + qk_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, + num_blocks: int, + block_write_idx: int, + eps: float, + output_norm_eps: float, +) -> torch.Tensor: + num_tokens, hidden_size = prefix.shape + assert prefix.stride(-1) == 1 + assert delta is None or delta.stride(-1) == 1 + assert blocks.stride(-1) == 1 + 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 + # The in-tree NVIDIA kernel covers the common fused-add + output-norm path; + # Triton handles block boundaries and final pre-norm output. + if ( + hidden_size == 7168 + and delta is not None + and output_norm_weight is not None + and num_blocks > 0 + and block_write_idx < 0 + and current_platform.is_device_capability_family(100) + ): + return ops.kimi_k3_attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + eps, + output_norm_eps, + ) + 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. + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + output, + prefix.stride(0), + 0 if delta is None else delta.stride(0), + blocks.stride(0), + blocks.stride(1), + output.stride(0), + num_blocks, + hidden_size, + block_write_idx, + eps, + output_norm_eps, + HAS_DELTA=delta is not None, + WRITE_BLOCK=block_write_idx >= 0, + APPLY_OUTPUT_NORM=output_norm_weight is not None, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(hidden_size), + num_warps=num_warps, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output From 5ed3faa43ddf075f24482396b634edc33e047a40 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 28 Jul 2026 15:46:51 +0800 Subject: [PATCH 14/67] [Rust Frontend] Add ordinary-text tokenizer encoding (#49992) Co-authored-by: OpenAI Codex Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/inkling/tests.rs | 4 + rust/src/parser/benches/utils/adapter.rs | 4 + rust/src/parser/src/unified/inkling.rs | 8 + rust/src/text/src/backend/hf/mod.rs | 4 + rust/src/tokenizer/src/hf.rs | 290 +++++++++++++++++++- rust/src/tokenizer/src/incremental.rs | 12 + rust/src/tokenizer/src/lib.rs | 4 + rust/src/tokenizer/src/tekken.rs | 54 ++++ rust/src/tokenizer/src/test_utils.rs | 30 ++ rust/src/tokenizer/src/tiktoken.rs | 40 +++ 10 files changed, 448 insertions(+), 2 deletions(-) diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index d0d55e1be2d..57e08ac0158 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -31,6 +31,10 @@ impl Tokenizer for FixtureTokenizer { Ok(text.bytes().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], diff --git a/rust/src/parser/benches/utils/adapter.rs b/rust/src/parser/benches/utils/adapter.rs index 243f19e3ce8..9cba325116c 100644 --- a/rust/src/parser/benches/utils/adapter.rs +++ b/rust/src/parser/benches/utils/adapter.rs @@ -19,6 +19,10 @@ impl Tokenizer for BenchTokenizer { Ok(text.chars().map(|_| u32::MAX).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs index eb78321d126..1ce69d8bfb4 100644 --- a/rust/src/parser/src/unified/inkling.rs +++ b/rust/src/parser/src/unified/inkling.rs @@ -414,6 +414,10 @@ mod tests { Ok(text.chars().map(u32::from).collect()) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, token_ids: &[u32], @@ -733,6 +737,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 49ae5dbd6b9..4fc7a18753a 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -177,6 +177,10 @@ mod tests { Ok(vec![]) } + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result> { + self.encode(text, false) + } + fn decode( &self, _token_ids: &[u32], diff --git a/rust/src/tokenizer/src/hf.rs b/rust/src/tokenizer/src/hf.rs index 08ec5a22d6b..eb527294a5b 100644 --- a/rust/src/tokenizer/src/hf.rs +++ b/rust/src/tokenizer/src/hf.rs @@ -1,13 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::borrow::Cow; use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use fastokens::Tokenizer as FastokensTokenizer; use fastokens::decoders::Decoder as FastokensDecoder; +use fastokens::pre_tokenized::{ + PreTokenizedString as FastokensPreTokenizedString, Split as FastokensSplit, +}; +use fastokens::{PreTokenizer as FastokensPreTokenizer, Split as FastokensSplitPreTokenizer}; use thiserror_ext::AsReport as _; -use tokenizers::Tokenizer as HfTokenizer; +use tokenizers::{ + AddedVocabulary, Model as _, OffsetType, PreTokenizer as _, Tokenizer as HfTokenizer, +}; use tracing::{info, warn}; use crate::byte_level_decode::decode_byte_level; @@ -16,6 +23,8 @@ use crate::{Result, Tokenizer}; mod added_tokens; +static EMPTY_HF_ADDED_VOCABULARY: LazyLock = LazyLock::new(AddedVocabulary::new); + enum Backend { Hf(Box), Fastokens(Box), @@ -53,6 +62,85 @@ fn decode_fastokens_byte_level( Ok(decode_byte_level(tokens)) } +fn encode_hf_ordinary(tokenizer: &HfTokenizer, text: &str) -> tokenizers::Result> { + let mut pretokenized = + EMPTY_HF_ADDED_VOCABULARY.extract_and_normalize(tokenizer.get_normalizer(), text); + + if let Some(pre_tokenizer) = tokenizer.get_pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized.tokenize(|normalized| tokenizer.get_model().tokenize(normalized.get()))?; + let encoding = pretokenized.into_encoding(None, 0, OffsetType::Byte)?; + let encoding = tokenizer.post_process(encoding, None, false)?; + Ok(encoding.get_ids().to_vec()) +} + +fn fastokens_fused_split(tokenizer: &FastokensTokenizer) -> Option<&FastokensSplitPreTokenizer> { + // Keep this predicate aligned with fastokens::Tokenizer::detect_fused_byte_level. + let FastokensPreTokenizer::Sequence(steps) = tokenizer.pre_tokenizer()? else { + return None; + }; + let [ + FastokensPreTokenizer::Split(split), + FastokensPreTokenizer::ByteLevel(byte_level), + ] = steps.as_slice() + else { + return None; + }; + byte_level.is_bulk_only().then_some(split) +} + +fn fastokens_pre_tokenized_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> FastokensPreTokenizedString { + // This is fastokens::Tokenizer::build_pre_tokenized with added_tokens = None. + let normalized = tokenizer + .normalizer() + .map_or(Cow::Borrowed(text), |normalizer| normalizer.normalize(text)); + match normalized { + Cow::Borrowed(_) => FastokensPreTokenizedString::from_text(text), + Cow::Owned(text) => { + let len = text.len(); + FastokensPreTokenizedString::new( + text, + vec![FastokensSplit { + range: 0..len, + token_id: None, + }], + ) + } + } +} + +fn encode_fastokens_ordinary( + tokenizer: &FastokensTokenizer, + text: &str, +) -> std::result::Result, fastokens::Error> { + if text.is_empty() { + return Ok(Vec::new()); + } + + let mut pretokenized = fastokens_pre_tokenized_ordinary(tokenizer, text); + let ids = if let Some(split) = fastokens_fused_split(tokenizer) { + split.pre_tokenize(&mut pretokenized)?; + pretokenized + .tokenize_batched(|buffer, splits, output| { + tokenizer.model().tokenize_batch_fused(buffer, splits, output) + }) + .map_err(fastokens::Error::Model)? + } else { + if let Some(pre_tokenizer) = tokenizer.pre_tokenizer() { + pre_tokenizer.pre_tokenize(&mut pretokenized)?; + } + pretokenized + .tokenize(|text, output| tokenizer.model().tokenize_into(text, output)) + .map_err(fastokens::Error::Model)? + }; + + Ok(tokenizer.post_process(ids, false)) +} + /// Tokenizer from `tokenizer.json` in HuggingFace format. /// /// This tries to load with `fastokens` first for better performance, then falls @@ -156,6 +244,17 @@ impl Tokenizer for HuggingFaceTokenizer { } } + fn encode_ordinary(&self, text: &str) -> Result> { + match &self.backend { + Backend::Hf(tokenizer) => encode_hf_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())), + Backend::Fastokens(tokenizer) | Backend::FastokensByteLevel(tokenizer) => { + encode_fastokens_ordinary(tokenizer, text) + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())) + } + } + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { match &self.backend { Backend::Hf(t) => t @@ -200,12 +299,19 @@ impl Tokenizer for HuggingFaceTokenizer { #[cfg(test)] mod tests { + use std::path::{Path, PathBuf}; + + use serde_json::{Value, json}; use tempfile::tempdir; use tokenizers::models::bpe::BPE; + use tokenizers::pre_tokenizers::byte_level::ByteLevel; use tokenizers::{AddedToken, Tokenizer as HfTokenizer}; use super::{HuggingFaceTokenizer, Tokenizer}; + const REGULAR_TOKEN: &str = "<|regular|>"; + const SPECIAL_TOKEN: &str = "<|special|>"; + fn tiny_bpe_tokenizer() -> HfTokenizer { let vocab = [ ("".to_string(), 0), @@ -232,6 +338,186 @@ mod tests { HfTokenizer::new(model) } + fn ordinary_test_tokenizer_json(fused: bool, with_added_tokens: bool) -> Value { + let mut alphabet: Vec = ByteLevel::alphabet().into_iter().collect(); + alphabet.sort_unstable(); + let vocab = alphabet + .into_iter() + .enumerate() + .map(|(id, token)| (token.to_string(), json!(id))) + .collect::>(); + + let pre_tokenizer = if fused { + json!({ + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": {"Regex": "\\S+|\\s+"}, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + }) + } else { + json!({ + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }) + }; + let added_tokens = with_added_tokens.then(|| { + json!([ + { + "id": 256, + "content": REGULAR_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": true, + "special": false + }, + { + "id": 257, + "content": SPECIAL_TOKEN, + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ]) + }); + + json!({ + "version": "1.0", + "truncation": { + "direction": "Right", + "max_length": 24, + "strategy": "LongestFirst", + "stride": 0 + }, + "padding": null, + "added_tokens": added_tokens.unwrap_or_else(|| json!([])), + "normalizer": {"type": "NFC"}, + "pre_tokenizer": pre_tokenizer, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": null, + "end_of_word_suffix": null, + "fuse_unk": false, + "byte_fallback": false, + "ignore_merges": false, + "vocab": vocab, + "merges": [] + } + }) + } + + fn write_tokenizer_json(dir: &Path, name: &str, value: &Value) -> PathBuf { + let path = dir.join(name); + std::fs::write( + &path, + serde_json::to_vec(value).expect("serialize tokenizer"), + ) + .expect("write tokenizer"); + path + } + + fn assert_ordinary_matches_added_empty( + constructor: fn(&Path) -> crate::Result, + fused: bool, + ) { + let dir = tempdir().expect("create temp dir"); + let added_path = write_tokenizer_json( + dir.path(), + "with-added.json", + &ordinary_test_tokenizer_json(fused, true), + ); + let empty_path = write_tokenizer_json( + dir.path(), + "added-empty.json", + &ordinary_test_tokenizer_json(fused, false), + ); + let tokenizer = constructor(&added_path).expect("load tokenizer with added tokens"); + let added_empty = constructor(&empty_path).expect("load tokenizer with empty added tokens"); + + if let super::Backend::Fastokens(inner) | super::Backend::FastokensByteLevel(inner) = + &tokenizer.backend + { + assert_eq!(super::fastokens_fused_split(inner).is_some(), fused); + } + + assert_eq!( + tokenizer.encode(REGULAR_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(REGULAR_TOKEN).unwrap()] + ); + assert_eq!( + tokenizer.encode(SPECIAL_TOKEN, false).unwrap(), + vec![tokenizer.token_to_id(SPECIAL_TOKEN).unwrap()] + ); + + for text in [ + "", + "hello", + "Cafe\u{301}", + REGULAR_TOKEN, + SPECIAL_TOKEN, + "hello <|regular|> Cafe\u{301} <|special|> tail", + ] { + assert_eq!( + tokenizer.encode_ordinary(text).unwrap(), + added_empty.encode(text, false).unwrap(), + "fused={fused}, text={text:?}", + ); + } + if matches!(&tokenizer.backend, super::Backend::Hf(_)) { + assert_eq!( + tokenizer + .encode_ordinary("hello <|regular|> Cafe\u{301} <|special|> tail") + .unwrap() + .len(), + 24, + "HF post-processing must retain configured truncation", + ); + } + } + + #[test] + fn hf_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_hf, fused); + } + } + + #[test] + fn fastokens_ordinary_matches_original_encode_with_added_empty() { + for fused in [false, true] { + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_fastokens, fused); + } + } + #[test] fn hf_constructor_resolves_added_token_ids() { let mut tokenizer = tiny_bpe_tokenizer(); diff --git a/rust/src/tokenizer/src/incremental.rs b/rust/src/tokenizer/src/incremental.rs index 5e470ae4b00..f608fa874b8 100644 --- a/rust/src/tokenizer/src/incremental.rs +++ b/rust/src/tokenizer/src/incremental.rs @@ -199,6 +199,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result { let bytes = token_ids.iter().map(|id| *id as u8).collect::>(); Ok(String::from_utf8_lossy(&bytes).into_owned()) @@ -273,6 +277,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let mut text = String::new(); for &token_id in token_ids { @@ -410,6 +418,10 @@ mod tests { unreachable!() } + fn encode_ordinary(&self, _text: &str) -> Result> { + unreachable!() + } + fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result { match token_ids { [1] => Ok("abc".into()), diff --git a/rust/src/tokenizer/src/lib.rs b/rust/src/tokenizer/src/lib.rs index 6c9fcd3fdea..0f8c7dc16e5 100644 --- a/rust/src/tokenizer/src/lib.rs +++ b/rust/src/tokenizer/src/lib.rs @@ -25,6 +25,10 @@ pub trait Tokenizer: Send + Sync { /// Encode one prompt string into token IDs. fn encode(&self, text: &str, add_special_tokens: bool) -> Result>; + /// Equivalent to `encode(text, false)`, except that every added, + /// special, and control-token matcher is bypassed. + fn encode_ordinary(&self, text: &str) -> Result>; + /// Decode one token sequence into text. fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result; diff --git a/rust/src/tokenizer/src/tekken.rs b/rust/src/tokenizer/src/tekken.rs index 20b6c26ffc8..5f9342e4a5e 100644 --- a/rust/src/tokenizer/src/tekken.rs +++ b/rust/src/tokenizer/src/tekken.rs @@ -35,6 +35,12 @@ impl Tokenizer for TekkenTokenizer { .map_err(|error| tokenizer_error!("encoding failed: {error}")) } + fn encode_ordinary(&self, text: &str) -> Result> { + self.inner + .encode(text, false, false) + .map_err(|error| tokenizer_error!("encoding failed: {error}")) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let policy = if skip_special_tokens { tekken::SpecialTokenPolicy::Ignore @@ -67,3 +73,51 @@ impl Tokenizer for TekkenTokenizer { self.inner.is_special_token(token_id) } } + +#[cfg(test)] +mod tests { + use base64::Engine as _; + use tekken::config::TokenizerVersion; + use tekken::{SpecialTokenInfo, TokenInfo}; + + use super::*; + + fn test_tokenizer() -> TekkenTokenizer { + let vocab = (0_u8..=255) + .map(|byte| TokenInfo { + rank: byte as usize, + token_bytes: base64::engine::general_purpose::STANDARD.encode([byte]), + token_str: None, + }) + .collect(); + let special_tokens = vec![SpecialTokenInfo { + rank: 0, + token_str: "".to_string(), + is_control: true, + }]; + let inner = Tekkenizer::new( + vocab, + &special_tokens, + r"(?s).", + 257, + 1, + TokenizerVersion::V3, + None, + ) + .expect("build Tekken tokenizer"); + TekkenTokenizer { inner } + } + + #[test] + fn ordinary_matches_tekkens_empty_special_encoding() { + let tokenizer = test_tokenizer(); + let text = "user text"; + let control_id = tokenizer.token_to_id("").unwrap(); + let ordinary_ids = tokenizer.encode_ordinary(text).unwrap(); + + assert_eq!(control_id, 0); + assert_eq!(ordinary_ids, tokenizer.encode(text, false).unwrap()); + assert!(!ordinary_ids.contains(&control_id)); + assert_eq!(tokenizer.decode(&ordinary_ids, false).unwrap(), text); + } +} diff --git a/rust/src/tokenizer/src/test_utils.rs b/rust/src/tokenizer/src/test_utils.rs index 36d1f3d18aa..6fdb6e02721 100644 --- a/rust/src/tokenizer/src/test_utils.rs +++ b/rust/src/tokenizer/src/test_utils.rs @@ -208,6 +208,10 @@ impl Tokenizer for TestTokenizer { Ok(ids) } + fn encode_ordinary(&self, text: &str) -> Result> { + Ok(text.as_bytes().iter().copied().map(u32::from).collect()) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { let mut output = String::new(); let mut pending_bytes = Vec::new(); @@ -374,6 +378,32 @@ mod tests { assert!(!tokenizer.is_special_id(0xF002)); } + #[test] + fn ordinary_encoding_bypasses_all_configured_tokens() { + let tokenizer = TestTokenizer::new() + .with_bos_token("", 256) + .with_special_token("", 257) + .with_regular_token("", 258); + let ordinary_text = "user and "; + + assert_eq!(tokenizer.encode("", false).unwrap(), vec![257]); + assert_eq!(tokenizer.encode("", false).unwrap(), vec![258]); + assert_eq!( + tokenizer.encode_ordinary(ordinary_text).unwrap(), + ordinary_text.as_bytes().iter().copied().map(u32::from).collect::>() + ); + + let mut segmented = tokenizer.encode("", false).unwrap(); + segmented.extend(tokenizer.encode_ordinary(ordinary_text).unwrap()); + segmented.extend(tokenizer.encode("", false).unwrap()); + assert_eq!(segmented.first(), Some(&257)); + assert_eq!(segmented.last(), Some(&258)); + assert_eq!( + tokenizer.decode(&segmented, false).unwrap(), + format!("{ordinary_text}") + ); + } + #[test] #[should_panic(expected = "configured test token id 255 overlaps byte fallback range 0..=255")] fn configured_token_id_must_stay_outside_byte_range() { diff --git a/rust/src/tokenizer/src/tiktoken.rs b/rust/src/tokenizer/src/tiktoken.rs index d5ca119b808..0f355566b94 100644 --- a/rust/src/tokenizer/src/tiktoken.rs +++ b/rust/src/tokenizer/src/tiktoken.rs @@ -462,6 +462,13 @@ impl Tokenizer for TiktokenTokenizer { }) } + fn encode_ordinary(&self, text: &str) -> Result> { + Ok(match &self.backend { + Backend::Riptoken(backend) => backend.inner.encode_ordinary(text), + Backend::TiktokenRs(backend) => backend.inner.encode_ordinary(text), + }) + } + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { // Filter passes: // @@ -752,6 +759,39 @@ mod tests { } } + #[test] + fn tiktoken_ordinary_bypasses_every_registered_added_token() { + let dir = tempfile::tempdir().expect("create temp dir"); + let bpe_path = write_synthetic_bpe_file(dir.path()); + fs::write( + dir.path().join("tokenizer_config.json"), + r#"{ + "added_tokens_decoder": { + "257": { "content": "<|im_end|>", "special": true }, + "258": { "content": "<|tool_call_begin|>", "special": false } + } + }"#, + ) + .expect("write tokenizer_config.json"); + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 260}"#) + .expect("write config.json"); + + let input = "<|im_end|><|tool_call_begin|><|reserved_token_259|>"; + let expected: Vec = input.as_bytes().iter().copied().map(u32::from).collect(); + for backend in explicit_backends(&bpe_path) { + assert_eq!(backend.encode("<|im_end|>", false).unwrap(), vec![257]); + assert_eq!( + backend.encode("<|tool_call_begin|>", false).unwrap(), + vec![258] + ); + assert_eq!( + backend.encode("<|reserved_token_259|>", false).unwrap(), + vec![259] + ); + assert_eq!(backend.encode_ordinary(input).unwrap(), expected); + } + } + /// `vocab_size` may live under `text_config` for composite (e.g. /// multimodal) configs. #[test] From 88402a41c4ab272ebbbd33f4a77fbbac0431cbb9 Mon Sep 17 00:00:00 2001 From: Liangliang Ma Date: Tue, 28 Jul 2026 16:26:49 +0800 Subject: [PATCH 15/67] [Test] Skip ROCm AITER MLA prefill tests on non-ROCm platforms (#49945) Signed-off-by: Liangliang-Ma Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_prefill_selector.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index f5985e7bc8e..e6d9f939ea5 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -8,6 +8,7 @@ import pytest import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig +from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum @@ -287,6 +288,11 @@ class TestBackendValidation: assert invalid_reasons == [] +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Imports vllm.platforms.rocm, whose module init requires a CUDA or " + "ROCm torch build; not importable on XPU/CPU/TPU.", +) class TestROCmAiterFAPrefillSelection: """Tests for the ROCm AITER FlashAttention MLA prefill backend.""" From 247470f23a4c8191a296c1993c4e86727ffb5191 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 28 Jul 2026 05:37:01 -0400 Subject: [PATCH 16/67] [CI] Add PyTorch stable ABI audit check (#48164) Signed-off-by: Chris Leonard Co-authored-by: Shengqi Chen --- .buildkite/check-torch-abi.py | 102 +++++++++++++++++++++++++++ .buildkite/ci_config.yaml | 1 + .buildkite/test_areas/torch_abi.yaml | 14 ++++ requirements/test/cpu.txt | 6 ++ requirements/test/cuda.in | 2 +- requirements/test/cuda.txt | 6 ++ 6 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 .buildkite/check-torch-abi.py create mode 100644 .buildkite/test_areas/torch_abi.yaml diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py new file mode 100644 index 00000000000..493952c33ec --- /dev/null +++ b/.buildkite/check-torch-abi.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Audit vLLM compiled libraries for PyTorch stable ABI compliance.""" + +import fnmatch +import sys +from pathlib import Path + +from torch_abi_audit import inspect_package +from torch_abi_audit.report import ExtensionReport, PackageReport + +# Temporary allowlist of extensions not yet on the stable ABI. +# Shrink and remove over time. +ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( + "vllm_flash_attn/_vllm_fa2_C.abi3.so", + "vllm_flash_attn/_vllm_fa3_C.abi3.so", + "third_party/deep_gemm/_C*.so", +) + + +def _relative_path(lib: ExtensionReport, package_root: Path) -> str: + try: + return lib.path.relative_to(package_root).as_posix() + except ValueError: + return lib.path.name + + +def _is_torch_unstable(lib: ExtensionReport) -> bool: + return lib.error is None and lib.torch.uses_torch and not lib.torch.stable + + +def _matches_allowlist(rel_path: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatch(rel_path, pattern) for pattern in patterns) + + +def _iter_libs(report: PackageReport) -> tuple[ExtensionReport, ...]: + return (*report.extensions, *report.bundled_libs) + + +def _collect_unstable(report: PackageReport) -> list[str]: + return sorted( + _relative_path(lib, report.root) + for lib in _iter_libs(report) + if _is_torch_unstable(lib) + ) + + +def _find_stale_allowlist_entries( + report: PackageReport, patterns: tuple[str, ...] +) -> list[str]: + """Allowlist patterns that match a built library which is no longer unstable.""" + stale: list[str] = [] + for pattern in patterns: + for lib in _iter_libs(report): + if lib.error is not None: + continue + if not fnmatch.fnmatch(_relative_path(lib, report.root), pattern): + continue + if not _is_torch_unstable(lib): + stale.append(pattern) + break + return stale + + +def check_torch_abi( + package: str = "vllm", + patterns: tuple[str, ...] = ALLOWED_UNSTABLE_LIBRARIES, +) -> int: + report = inspect_package(package) + if report.error: + print(f"error: failed to inspect {package!r}: {report.error}", file=sys.stderr) + return 2 + + unstable = _collect_unstable(report) + unexpected = [ + rel_path for rel_path in unstable if not _matches_allowlist(rel_path, patterns) + ] + stale = _find_stale_allowlist_entries(report, patterns) + + if unexpected or stale: + if unexpected: + print( + "Not allowed: torch-unstable libraries outside " + f"ALLOWED_UNSTABLE_LIBRARIES: {', '.join(unexpected)}", + file=sys.stderr, + ) + if stale: + print( + "Not allowed: stale ALLOWED_UNSTABLE_LIBRARIES entries: " + f"{', '.join(stale)}", + file=sys.stderr, + ) + return 1 + + print("Torch stable ABI check passed.") + return 0 + + +if __name__ == "__main__": + print(">>> Auditing vLLM extension modules for PyTorch stable ABI compliance") + sys.exit(check_torch_abi()) diff --git a/.buildkite/ci_config.yaml b/.buildkite/ci_config.yaml index 21ffa1b9b8d..9e1e46db67e 100644 --- a/.buildkite/ci_config.yaml +++ b/.buildkite/ci_config.yaml @@ -14,6 +14,7 @@ run_all_patterns: - "setup.py" - "csrc/" - "cmake/" + - ".buildkite/check-torch-abi.py" run_all_exclude_patterns: - "docker/Dockerfile." - "csrc/cpu/" diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml new file mode 100644 index 00000000000..eaef3551664 --- /dev/null +++ b/.buildkite/test_areas/torch_abi.yaml @@ -0,0 +1,14 @@ +group: Torch ABI +depends_on: + - image-build +steps: +- label: Torch Stable ABI Audit + key: torch-stable-abi-audit + timeout_in_minutes: 5 + source_file_dependencies: + - .buildkite/check-torch-abi.py + - csrc/ + - cmake/ + - setup.py + commands: + - python3 /vllm-workspace/.buildkite/check-torch-abi.py diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 3cb251a308b..923b54bb14e 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -763,6 +765,8 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -r requirements/test/../common.txt @@ -1127,6 +1131,8 @@ torch==2.13.0+cpu # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cpu # via # -r requirements/test/cuda.in diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index b33257250e7..377224eac36 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -45,7 +45,7 @@ schemathesis>=4.0.0 # Required for openai schema test. # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 - +torch-abi-audit # CI check for PyTorch stable ABI compliance genai_perf>=0.0.8 tritonclient>=2.51.0 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 6490502cdda..6c155d89bd2 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1,5 +1,7 @@ # This file was autogenerated by uv via the following command: # uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12 +abi3info==2025.11.29 + # via torch-abi-audit absl-py==2.1.0 # via rouge-score accelerate==1.13.0 @@ -850,6 +852,8 @@ pycparser==2.22 # via cffi pycryptodomex==3.22.0 # via blobfile +pycxxfilt==0.1.0 + # via torch-abi-audit pydantic==2.12.0 # via # -c requirements/common.txt @@ -1225,6 +1229,8 @@ torch==2.13.0+cu130 # vector-quantize-pytorch # vocos # xgrammar +torch-abi-audit==0.0.1 + # via -r requirements/test/cuda.in torchaudio==2.11.0+cu130 # via # -c requirements/cuda.txt From 25ace8fe5df07fc13f4aef5a89db391f326e60ee Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 28 Jul 2026 18:04:36 +0800 Subject: [PATCH 17/67] [CI] Increase Qwen3.5 MTP GSM8K generation length (#49881) Signed-off-by: zjy0516 Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml index d247515a0f0..921365ae686 100644 --- a/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml +++ b/tests/evals/gsm8k/configs/Qwen3.5-397B-A17B-NVFP4-DEP2-MTP.yaml @@ -3,8 +3,9 @@ accuracy_threshold: 0.88 tolerance: 0.03 num_questions: 1319 num_fewshot: 5 +max_tokens: 12000 server_args: >- - --max-model-len 4096 + --max-model-len 16384 --data-parallel-size 2 --enable-expert-parallel --max-num-seqs 384 From bf9f23003c67698b973a674b54e4ca77aae1adff Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:53:18 +0800 Subject: [PATCH 18/67] [Rust Frontend] Fix finish reason for named tool choices (#49496) Signed-off-by: reidliu41 --- .../src/routes/openai/chat_completions.rs | 16 ++-- .../routes/openai/chat_completions/convert.rs | 6 ++ rust/src/server/src/routes/tests.rs | 83 ++++++++++++++++--- 3 files changed, 88 insertions(+), 17 deletions(-) diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index bd86a06136a..8a1e9c4383a 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -131,6 +131,7 @@ async fn collect_chat_completion( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, ) -> Result { let collected = stream.collect_message().await.map_err(|error| { @@ -157,7 +158,9 @@ async fn collect_chat_completion( // When reasoning is hidden, omit them rather than leaking hidden reasoning // tokens through per-token metadata. let include_output_metadata = include_reasoning || reasoning.is_none(); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?.to_string(); + let finish_reason = + chat_finish_reason_to_openai(&finish_reason, saw_tool_calls && !is_named_tool_choice)? + .to_string(); let tool_calls = message .tool_calls() .map(|call| ToolCall { @@ -254,6 +257,7 @@ async fn chat_completion_chunk_stream( echo, return_token_ids, return_tokens_as_token_ids, + is_named_tool_choice, }: ResponseOptions, mut y: TryYielder, ) -> Result<(), ApiError> { @@ -454,7 +458,7 @@ async fn chat_completion_chunk_stream( &response_model, created, finish_reason, - saw_tool_calls, + saw_tool_calls && !is_named_tool_choice, ) { Ok(chunk) => yield_chunk!(chunk), Err(error) => { @@ -787,10 +791,10 @@ fn final_chunk( response_model: &str, created: u64, finish_reason: FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result { let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); - let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?; + let finish_reason = chat_finish_reason_to_openai(&finish_reason, use_tool_calls_finish_reason)?; debug!( finish_reason = %finish_reason, @@ -809,10 +813,10 @@ fn final_chunk( fn chat_finish_reason_to_openai( finish_reason: &FinishReason, - saw_tool_calls: bool, + use_tool_calls_finish_reason: bool, ) -> Result<&'static str, ApiError> { match finish_reason { - FinishReason::Stop(_) if saw_tool_calls => Ok("tool_calls"), + FinishReason::Stop(_) if use_tool_calls_finish_reason => Ok("tool_calls"), FinishReason::Stop(_) => Ok("stop"), FinishReason::Length => Ok("length"), FinishReason::Abort => Ok("abort"), diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 5b2fa3c19ed..d75f3fab162 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -54,6 +54,8 @@ pub(super) struct ResponseOptions { pub return_token_ids: bool, /// Whether to format logprob tokens as `token_id:{id}`. pub return_tokens_as_token_ids: bool, + /// Whether the request forces one named function tool. + pub is_named_tool_choice: bool, } /// Validate and lower one OpenAI chat completion request into the internal chat @@ -98,6 +100,7 @@ pub(super) fn prepare_chat_request( .and_then(|options| options.continuous_usage_stats) .unwrap_or(false); let requested_logprobs = request.logprobs; + let is_named_tool_choice = matches!(&request.tool_choice, Some(ToolChoice::Function { .. })); // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's // behavior. @@ -180,6 +183,7 @@ pub(super) fn prepare_chat_request( echo, return_token_ids: request.return_token_ids.unwrap_or(false), return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false), + is_named_tool_choice, }, chat_request, }) @@ -1068,6 +1072,7 @@ mod tests { .expect("request is valid"); assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Required); + assert!(!prepared.options.is_named_tool_choice); } #[test] @@ -1107,6 +1112,7 @@ mod tests { name: "get_weather".to_string(), } ); + assert!(prepared.options.is_named_tool_choice); } #[test] diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 64e904b562b..1c228f99f1f 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -153,6 +153,20 @@ fn default_stream_output_specs() -> Vec<(Vec, Option Vec<(Vec, Option)> { + vec![ + (bytes_to_token_ids(b"Need tool."), None), + ( + bytes_to_token_ids(b"\n{\"name\":\"get_weather\", "), + None, + ), + ( + bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n"), + Some(EngineCoreFinishReason::Stop), + ), + ] +} + fn assert_adapter_a_lora_request(request: &EngineCoreRequest) { let lora = request.lora_request.as_ref().expect("lora request"); assert_eq!(lora.lora_name, "adapter-a"); @@ -4554,17 +4568,7 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() { async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { let (app, engine_task) = test_app_with_backend_and_stream_output_specs( Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), - vec![ - (bytes_to_token_ids(b"Need tool."), None), - ( - bytes_to_token_ids(b"\n{\"name\":\"get_weather\", "), - None, - ), - ( - bytes_to_token_ids(b"\"arguments\":{\"city\":\"Paris\"}}\n"), - Some(EngineCoreFinishReason::Stop), - ), - ], + weather_tool_call_output_specs(), ) .await; @@ -4613,6 +4617,63 @@ async fn tool_calls_are_mapped_to_tool_call_sse_chunks() { assert!(text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn named_tool_choice_uses_stop_finish_reason() { + for stream in [false, true] { + let (app, engine_task) = test_app_with_backend_and_stream_output_specs( + Arc::new(FakeChatBackend::with_model_id("Qwen/Qwen3-0.6B")), + weather_tool_call_output_specs(), + ) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": stream, + "messages": [{"role": "user", "content": "hello"}], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}} + } + } + }], + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"} + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + + assert!(text.contains("\"tool_calls\":"), "{text}"); + assert!(text.contains("\"name\":\"get_weather\""), "{text}"); + assert!(text.contains("\"finish_reason\":\"stop\""), "{text}"); + assert!(!text.contains("\"finish_reason\":\"tool_calls\""), "{text}"); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn tool_call_sse_chunks_can_carry_logprobs() { From 912d6b619de0f2a44df74704973cff90fa0fb4e9 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:02:00 +0800 Subject: [PATCH 19/67] [Rust Frontend] Align sampling validation with Python (#47494) Signed-off-by: reidliu41 --- rust/src/server/src/error.rs | 16 ++++ rust/src/server/src/grpc/tests.rs | 29 +++++++ rust/src/text/src/error.rs | 4 + rust/src/text/src/lib.rs | 2 +- rust/src/text/src/lower.rs | 119 +++++++++++++++++++++++++++- rust/src/text/src/lower/sampling.rs | 92 +++++++++++++++++++++ 6 files changed, 260 insertions(+), 2 deletions(-) create mode 100644 rust/src/text/src/lower/sampling.rs diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index fbd13b77baf..3a8c0ae5c92 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -141,6 +141,22 @@ mod tests { assert!(response.error.message.contains("max_tokens=4")); } + #[test] + fn sampling_params_validation_maps_to_invalid_request() { + let api_error = text_submit_error( + "failed to submit completion request", + vllm_text::Error::SamplingParams(vllm_text::SamplingParamsError::OutOfRange { + parameter: "top_p", + value: 0.0, + expected: "(0, 1]", + }), + ); + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST); + let response = api_error.to_error_response(); + assert_eq!(response.error.error_type, "invalid_request_error"); + assert!(response.error.message.contains("top_p")); + } + #[test] fn chat_wrapped_prompt_too_long_maps_to_invalid_request() { let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 75da670ec3b..808f3422758 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -648,6 +648,35 @@ async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = grpc_test_server( + b"engine-grpc-invalid-sampling", + default_stream_output_specs(), + ) + .await; + + let status = client + .generate(pb::GenerateRequest { + request_id: "test-invalid-sampling".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), + sampling: Some(pb::RandomSampling { + top_p: 2.0, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect_err("should fail when top_p is out of range"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("top_p")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn streaming_generate_yields_incremental_responses() { diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index e72e68196c3..6b385fad68d 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -6,6 +6,7 @@ use vllm_engine_core_client::Error as EngineCoreError; use vllm_llm::Error as LlmError; pub use crate::lower::logprobs::LogprobsError; +pub use crate::lower::sampling::SamplingParamsError; pub use crate::lower::token_ids::TokenIdsError; #[derive(Debug, Error)] @@ -23,6 +24,8 @@ pub enum Error { Logprobs(#[from] LogprobsError), #[error(transparent)] TokenIds(#[from] TokenIdsError), + #[error(transparent)] + SamplingParams(#[from] SamplingParamsError), #[error( "`min_tokens` must be less than or equal to `max_tokens`, \ got min_tokens={min_tokens}, max_tokens={max_tokens}" @@ -50,6 +53,7 @@ impl Error { | Self::EmptyPromptTokenIds { .. } | Self::Logprobs(_) | Self::TokenIds(_) + | Self::SamplingParams(_) | Self::MinTokensExceedsMaxTokens { .. } | Self::InvalidThinkingTokenBudget | Self::InvalidRepetitionDetection { .. } diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index b00155999e8..646e4bac9f3 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -10,7 +10,7 @@ use std::mem::take; pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend}; -pub use error::{Error, LogprobsError, Result, TokenIdsError}; +pub use error::{Error, LogprobsError, Result, SamplingParamsError, TokenIdsError}; use futures::Stream; pub use lower::{ PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index bd43a1d141c..aa54595acda 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -4,9 +4,11 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; +pub(crate) mod sampling; pub(crate) mod token_ids; use logprobs::validate_logprobs; +use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ EngineCoreSamplingParams, RepetitionDetectionParams, @@ -186,6 +188,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, }; + validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; Ok(params) } @@ -319,7 +322,7 @@ mod tests { use super::*; use crate::backend::hf::HfTextBackend; use crate::backend::{SamplingHints, TextBackend as _}; - use crate::error::{LogprobsError, TokenIdsError}; + use crate::error::{LogprobsError, SamplingParamsError, TokenIdsError}; use crate::request::{Prompt, TextRequest}; fn stub_tokenizer() -> TestTokenizer { @@ -482,6 +485,120 @@ mod tests { assert!(message.contains("min_count=1")); } + #[test] + fn lower_sampling_params_rejects_invalid_sampling_ranges() { + let cases = [ + ( + "temperature", + SamplingParams { + temperature: Some(5.0), + ..SamplingParams::default() + }, + ), + ( + "top_p", + SamplingParams { + top_p: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "min_p", + SamplingParams { + min_p: Some(2.0), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(0.0), + ..SamplingParams::default() + }, + ), + ( + "frequency_penalty", + SamplingParams { + frequency_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ( + "presence_penalty", + SamplingParams { + presence_penalty: Some(100.0), + ..SamplingParams::default() + }, + ), + ]; + + for (expected_parameter, sampling_params) in cases { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::OutOfRange { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should be rejected" + ); + } + } + + #[test] + fn lower_sampling_params_rejects_non_finite_sampling_values() { + for (expected_parameter, sampling_params) in [ + ( + "temperature", + SamplingParams { + temperature: Some(f32::INFINITY), + ..SamplingParams::default() + }, + ), + ( + "repetition_penalty", + SamplingParams { + repetition_penalty: Some(f32::NAN), + ..SamplingParams::default() + }, + ), + ] { + let error = + lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) + .unwrap_err(); + + assert!( + matches!( + error, + Error::SamplingParams(SamplingParamsError::NotFinite { + parameter, + .. + }) if parameter == expected_parameter + ), + "{expected_parameter} should reject non-finite values" + ); + } + } + + #[test] + fn lower_sampling_params_accepts_python_compatible_repetition_penalty_above_two() { + let params = lower_sampling_params_with_limits( + SamplingParams { + repetition_penalty: Some(2.5), + ..SamplingParams::default() + }, + sample_sampling_limits(), + ) + .unwrap(); + + assert_eq!(params.repetition_penalty, 2.5); + } + #[test] fn lower_text_request_applies_python_style_eos_hints() { let prepared = lower_text_request( diff --git a/rust/src/text/src/lower/sampling.rs b/rust/src/text/src/lower/sampling.rs new file mode 100644 index 00000000000..edcdc4b0492 --- /dev/null +++ b/rust/src/text/src/lower/sampling.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use thiserror::Error; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; + +#[derive(Debug, Error, PartialEq)] +pub enum SamplingParamsError { + #[error("{parameter} must be a finite number, got {value}")] + NotFinite { parameter: &'static str, value: f32 }, + #[error("{parameter} must be in {expected}, got {value}")] + OutOfRange { + parameter: &'static str, + value: f32, + expected: &'static str, + }, +} + +fn validate_frequency_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("frequency_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_presence_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("presence_penalty", value, -2.0, 2.0, "[-2, 2]") +} + +fn validate_temperature(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("temperature", value)?; + validate_closed_range("temperature", value, 0.0, 2.0, "[0, 2]") +} + +fn validate_top_p(value: f32) -> Result<(), SamplingParamsError> { + if value > 0.0 && value <= 1.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "top_p", + value, + expected: "(0, 1]", + }) +} + +fn validate_min_p(value: f32) -> Result<(), SamplingParamsError> { + validate_closed_range("min_p", value, 0.0, 1.0, "[0, 1]") +} + +fn validate_repetition_penalty(value: f32) -> Result<(), SamplingParamsError> { + validate_finite("repetition_penalty", value)?; + if value > 0.0 { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter: "repetition_penalty", + value, + expected: "(0, inf)", + }) +} + +pub(crate) fn validate_resolved_sampling_params( + params: &EngineCoreSamplingParams, +) -> Result<(), SamplingParamsError> { + validate_temperature(params.temperature)?; + validate_top_p(params.top_p)?; + validate_min_p(params.min_p)?; + validate_frequency_penalty(params.frequency_penalty)?; + validate_presence_penalty(params.presence_penalty)?; + validate_repetition_penalty(params.repetition_penalty) +} + +fn validate_finite(parameter: &'static str, value: f32) -> Result<(), SamplingParamsError> { + if value.is_finite() { + return Ok(()); + } + Err(SamplingParamsError::NotFinite { parameter, value }) +} + +fn validate_closed_range( + parameter: &'static str, + value: f32, + min: f32, + max: f32, + expected: &'static str, +) -> Result<(), SamplingParamsError> { + if value >= min && value <= max { + return Ok(()); + } + Err(SamplingParamsError::OutOfRange { + parameter, + value, + expected, + }) +} From d2bfc6fe20343c638840c8867c29f7365fe23378 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 28 Jul 2026 05:09:19 -0700 Subject: [PATCH 20/67] [Build] Fix DeepEP CUDA driver stub linking (#50103) Signed-off-by: khluu Co-authored-by: OpenAI Codex --- tools/ep_kernels/install_python_libraries.sh | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tools/ep_kernels/install_python_libraries.sh b/tools/ep_kernels/install_python_libraries.sh index 739f031c9ef..5f5a597baca 100755 --- a/tools/ep_kernels/install_python_libraries.sh +++ b/tools/ep_kernels/install_python_libraries.sh @@ -197,6 +197,21 @@ do_build() { #endif' csrc/kernels/backend/symmetric.hpp fi + if [[ "$name" == "DeepEP" ]]; then + # DeepEP links against the CUDA driver API in driverless build images. + local cuda_driver_stub + local cuda_driver_stub_dir + cuda_driver_stub=$( + find -H "$CUDA_HOME" -path "*/stubs/libcuda.so" -print -quit + ) + if [[ -z "$cuda_driver_stub" ]]; then + echo "CUDA driver stub not found under $CUDA_HOME" >&2 + exit 1 + fi + cuda_driver_stub_dir=$(dirname "$cuda_driver_stub") + export LIBRARY_PATH="${cuda_driver_stub_dir}${LIBRARY_PATH:+:$LIBRARY_PATH}" + fi + if [ "$MODE" = "install" ]; then echo "Installing $name into environment" eval "$extra_env" uv pip install --no-build-isolation -vvv . From 35efdf6b34f1ef76c23c8980f6e8a3b73cab50f1 Mon Sep 17 00:00:00 2001 From: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:19:13 +0300 Subject: [PATCH 21/67] [Elastic EP] Async preparation (#47288) Signed-off-by: Itay Alroy --- .buildkite/test_areas/expert_parallelism.yaml | 1 + tests/distributed/test_elastic_ep.py | 166 +++++-- vllm/config/parallel.py | 9 + .../base_device_communicator.py | 16 +- .../device_communicators/cpu_communicator.py | 5 +- .../device_communicators/cuda_communicator.py | 2 + .../device_communicators/xpu_communicator.py | 5 +- .../distributed/elastic_ep/elastic_execute.py | 181 +++---- vllm/distributed/elastic_ep/elastic_state.py | 460 ++++++------------ vllm/distributed/elastic_ep/standby_state.py | 3 +- vllm/distributed/parallel_state.py | 14 +- vllm/distributed/stateless_coordinator.py | 2 + .../serve/elastic_ep/api_router.py | 9 +- vllm/model_executor/warmup/kernel_warmup.py | 31 +- vllm/v1/engine/__init__.py | 2 - vllm/v1/engine/async_llm.py | 45 +- vllm/v1/engine/coordinator.py | 3 + vllm/v1/engine/core.py | 75 ++- vllm/v1/engine/core_client.py | 310 ++++++------ vllm/v1/engine/utils.py | 31 +- vllm/v1/executor/abstract.py | 8 +- 21 files changed, 701 insertions(+), 677 deletions(-) diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index a3b46b58285..1d1609d46b8 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -52,4 +52,5 @@ steps: - vllm/compilation/ - tests/distributed/ commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - pytest -v -s distributed/test_elastic_ep.py diff --git a/tests/distributed/test_elastic_ep.py b/tests/distributed/test_elastic_ep.py index 4ce7497598a..01c254c2d03 100644 --- a/tests/distributed/test_elastic_ep.py +++ b/tests/distributed/test_elastic_ep.py @@ -3,7 +3,9 @@ import os import subprocess +import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest import requests @@ -40,6 +42,106 @@ def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool: return False +def _traffic_loop( + server: RemoteOpenAIServer, + dp_rank: int | None, + ready: threading.Barrier, + stop: threading.Event, + finished: threading.Event, + is_probe: bool = False, +) -> list[tuple[float, float, int | None]]: + url = server.url_for("is_scaling_elastic_ep" if is_probe else "v1/completions") + payload = {"model": MODEL_NAME, "prompt": "Hello", "max_tokens": 4} + headers = None if dp_rank is None else {"X-data-parallel-rank": str(dp_rank)} + request_payload = None if is_probe else payload + responses = [] + is_ready = False + while not stop.is_set(): + request_start = time.perf_counter() + try: + response = requests.post( + url, json=request_payload, headers=headers, timeout=120 + ) + status_code = response.status_code + except requests.exceptions.RequestException: + status_code = None + responses.append((request_start, time.perf_counter(), status_code)) + if status_code == 200: + if not is_ready: + ready.wait(timeout=120) + is_ready = True + if finished.is_set(): + return responses + time.sleep(0.05) + return responses + + +def _downtime(responses: list[tuple[float, float, int | None]]) -> float: + rejected = [end for _, end, status in responses if status == 503] + if not rejected: + return 0 + recovered = next( + end for _, end, status in responses if status == 200 and end > rejected[-1] + ) + return recovered - rejected[0] + + +def _scale_with_traffic( + server: RemoteOpenAIServer, + source_dp_size: int, + new_dp_size: int, + traffic_mode: str, +) -> None: + traffic_clients: list[int | None] = [] + if traffic_mode == "light": + traffic_clients = [0] + elif traffic_mode == "heavy": + traffic_clients = [None] * source_dp_size + clients = [(None, True)] + [(rank, False) for rank in traffic_clients] + ready = threading.Barrier(len(clients) + 1) + stop = threading.Event() + finished = threading.Event() + + with ThreadPoolExecutor(max_workers=len(clients)) as executor: + futures = [ + executor.submit( + _traffic_loop, server, rank, ready, stop, finished, is_probe + ) + for rank, is_probe in clients + ] + try: + ready.wait(timeout=120) + start_time = time.perf_counter() + assert _send_scale_command(server, new_dp_size) + scale_seconds = time.perf_counter() - start_time + finished.set() + probe_result, *results = [future.result(timeout=120) for future in futures] + finally: + stop.set() + + bad_statuses = { + status + for responses in [probe_result, *results] + for _, _, status in responses + if status not in (200, 503) + } + assert not bad_statuses, f"traffic got unexpected statuses {bad_statuses}" + probe_503 = [start for start, _, status in probe_result if status == 503] + assert probe_503, "Scaling probe did not observe commit" + assert not results or any( + status == 200 and start_time <= request_start and request_end < probe_503[0] + for responses in results + for request_start, request_end, status in responses + ), "No request completed successfully during preparation" + + print( + f"[Elastic EP timing][{source_dp_size}->{new_dp_size}]" + f"[traffic={traffic_mode}] " + f"scale_seconds={scale_seconds:.3f} " + f"downtime_seconds={_downtime(probe_result):.3f}" + ) + + def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: assert server.port is not None result = evaluate_gsm8k( @@ -59,7 +161,7 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float: return accuracy -def _base_serve_args(use_async_eplb: bool = False) -> list[str]: +def _base_serve_args(dp_size: int = 2, enforce_eager: bool = False) -> list[str]: args = [ "--trust-remote-code", "--tensor-parallel-size", @@ -78,57 +180,65 @@ def _base_serve_args(use_async_eplb: bool = False) -> list[str]: "--eplb-config.num_redundant_experts", "0", "--eplb-config.use_async", - "true" if use_async_eplb else "false", + "true", "--eplb-config.step_interval", - "10", + "300", "--eplb-config.window_size", "5", "--data-parallel-backend", "ray", "--data-parallel-size", - "2", + str(dp_size), "--api-server-count", "1", + "--disable-access-log-for-endpoints", + "/is_scaling_elastic_ep", ] leader_address = os.environ.get("LEADER_ADDRESS") if leader_address: args.extend(["--data-parallel-address", leader_address]) + if enforce_eager: + args.append("--enforce-eager") return args @pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] + ("enforce_eager", "traffic_mode"), + [ + pytest.param(True, "none", id="enforce_eager_none"), + pytest.param(True, "light", id="enforce_eager_light"), + pytest.param(True, "heavy", id="enforce_eager_heavy"), + pytest.param(False, "heavy", id="cuda_graphs_heavy"), + ], ) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling(use_async_eplb: bool): - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl +def test_elastic_ep_scaling(enforce_eager: bool, traffic_mode: str): + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + initial_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_INITIAL_DP", "2")) + target_dp_size = int(os.getenv("VLLM_TEST_ELASTIC_EP_TARGET_DP", "4")) + assert target_dp_size > initial_dp_size + vllm_serve_args = _base_serve_args(initial_dp_size, enforce_eager) with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 ) as server: - initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)") - - assert _send_scale_command(server, 4) - time.sleep(10) - scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)") + initial_accuracy = _run_gsm8k_eval(server, "Initial") + _scale_with_traffic(server, initial_dp_size, target_dp_size, traffic_mode) + scale_up_accuracy = _run_gsm8k_eval(server, "After scale up") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" ) - assert _send_scale_command(server, 2) - time.sleep(5) - scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") - + _scale_with_traffic(server, target_dp_size, initial_dp_size, traffic_mode) + scale_down_accuracy = _run_gsm8k_eval(server, "After scale down") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than " f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}" @@ -147,24 +257,20 @@ def test_elastic_ep_scaling(use_async_eplb: bool): print(f" Tolerance: {ACCURACY_TOL:.3f}") -@pytest.mark.parametrize( - "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"] -) @multi_gpu_test(num_gpus=4) -def test_elastic_ep_scaling_uneven(use_async_eplb: bool): +def test_elastic_ep_scaling_uneven(): """Test scale up with uneven worker distribution. This tests the case where num_new_workers % old_dp_size != 0, specifically 2 -> 3 where remainder = 1 % 2 = 1. This exercises the remainder handling in sender-receiver pairing. """ - if use_async_eplb: - from vllm.distributed.eplb.eplb_communicator import has_nixl + from vllm.distributed.eplb.eplb_communicator import has_nixl - if not has_nixl(): - pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") + if not has_nixl(): + pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)") - vllm_serve_args = _base_serve_args(use_async_eplb) + vllm_serve_args = _base_serve_args() with RemoteOpenAIServer( MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200 @@ -174,7 +280,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale 2 -> 3: This has remainder = 1 % 2 = 1 # Tests uneven sender-receiver pairing assert _send_scale_command(server, 3) - time.sleep(10) scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)") assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, ( @@ -184,7 +289,6 @@ def test_elastic_ep_scaling_uneven(use_async_eplb: bool): # Scale back down to 2 assert _send_scale_command(server, 2) - time.sleep(5) scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)") assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, ( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 949eb298a17..5ebc410f3c6 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -686,6 +686,14 @@ class ParallelConfig: and self.data_parallel_size > 1 ) + @property + def use_all2all(self) -> bool: + return ( + self.data_parallel_size > 1 + or self.use_sequence_parallel_moe + or (self.enable_expert_parallel and self.prefill_context_parallel_size > 1) + ) + @property def use_batched_dp_moe(self) -> bool: return ( @@ -786,6 +794,7 @@ class ParallelConfig: "data_parallel_master_ip", "data_parallel_master_port", "_data_parallel_master_port_list", + "_coord_store_port", "data_parallel_rpc_port", "rank", "master_addr", diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 45438a54691..73fd1331f5c 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -175,6 +175,7 @@ class DeviceCommunicatorBase: unique_name: str = "", global_ranks: list[int] | None = None, global_world_size: int | None = None, + use_all2all: bool = False, ): self.device = device or torch.device("cpu") self.cpu_group = cpu_group @@ -204,26 +205,15 @@ class DeviceCommunicatorBase: self.global_world_size = dist.get_world_size() self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank) - use_ep = False all2all_backend = None from vllm.config import get_current_vllm_config_or_none config = get_current_vllm_config_or_none() if config is not None: - # initialize the all2all manager for DP or sequence-parallel EP. - parallel_config = config.parallel_config - use_ep = ( - parallel_config.data_parallel_size > 1 - or parallel_config.use_sequence_parallel_moe - or ( - parallel_config.enable_expert_parallel - and parallel_config.prefill_context_parallel_size > 1 - ) - ) - all2all_backend = parallel_config.all2all_backend + all2all_backend = config.parallel_config.all2all_backend self.is_ep_communicator = unique_name.split(":")[0] == "ep" - self.use_all2all = self.is_ep_communicator and use_ep + self.use_all2all = self.is_ep_communicator and use_all2all self.all2all_backend = all2all_backend self.all2all_manager: All2AllManagerBase | None = None diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 9ec4b72f80d..8ea12d9255a 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -24,8 +24,11 @@ class CpuCommunicator(DeviceCommunicatorBase): device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.dist_module = torch.distributed if ( diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fccc6ba60c3..23e37ca830e 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -36,6 +36,7 @@ class CudaCommunicator(DeviceCommunicatorBase): global_ranks: list[int] | None = None, global_world_size: int | None = None, tcp_store_group: StatelessProcessGroup | None = None, + use_all2all: bool = False, ): super().__init__( cpu_group, @@ -44,6 +45,7 @@ class CudaCommunicator(DeviceCommunicatorBase): unique_name, global_ranks, global_world_size, + use_all2all=use_all2all, ) if "tp" not in unique_name: # custom allreduce or torch symm mem can be used only by tp diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py index 1b6ce9e8aae..7ca132824ec 100644 --- a/vllm/distributed/device_communicators/xpu_communicator.py +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -20,8 +20,11 @@ class XpuCommunicator(DeviceCommunicatorBase): device: torch.device | None = None, device_group: ProcessGroup | None = None, unique_name: str = "", + use_all2all: bool = False, ): - super().__init__(cpu_group, device, device_group, unique_name) + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) self.ca_comm: None = None if self.use_all2all: if self.all2all_backend in ("naive", "allgather_reducescatter"): diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index b0c3740f57e..cea7fcb2f01 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy import gc import weakref from collections.abc import Iterable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import replace from typing import TYPE_CHECKING @@ -43,6 +43,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.fused_moe.eep_reconfigure import ( make_eep_staged_quant_method, ) +from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.utils import is_moe_layer from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper @@ -145,6 +146,10 @@ class ElasticEPScalingExecutor: self.worker_ref = weakref.ref(worker) self.reconfig_request = None self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {} + self._async_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPAsync" + ) + self._async_future: Future[None] | None = None @property def worker(self): @@ -159,59 +164,68 @@ class ElasticEPScalingExecutor: raise ValueError(f"Unknown execute method: {execute_method}") return method(*args, **kwargs) - def _set_eplb_suppressed(self, suppressed: bool) -> None: - self.worker.model_runner.eep_eplb_suppressed = suppressed - ep_group = get_standby_ep_group() or get_ep_group() - if ep_group.rank == 0: - logger.info( - "[Elastic EP] EPLB %s elastic scaling transition", - "disabled during" if suppressed else "re-enabled after", - ) + def start_async(self, execute_method: str, *args, **kwargs) -> str: + if self._async_future is not None: + raise RuntimeError("Another Elastic EP async method is active") + if args and isinstance(args[0], ReconfigureDistributedRequest): + self.reconfig_request = args[0] + dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank + done_key = f"eep_async/{execute_method}/{dp_rank}/{self.worker.rank}" + self._async_future = self._async_executor.submit( + self._run_async, execute_method, *args, **kwargs + ) + self._async_future.add_done_callback(lambda _: self._mark_async_done(done_key)) + return done_key + + def _run_async(self, execute_method: str, *args, **kwargs) -> None: + from vllm.platforms import current_platform + + self.worker.vllm_config.enable_trace_function_call_for_thread() + assert hasattr(self.worker, "device") + current_platform.set_device(self.worker.device) + with set_current_vllm_config(self.worker.vllm_config): + self.execute(execute_method, *args, **kwargs) + + def _mark_async_done(self, done_key: str) -> None: + from vllm.distributed.utils import get_cached_tcp_store_client + + assert self.reconfig_request is not None + get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ).set(done_key, b"1") + + def clear_async(self) -> None: + future = self._async_future + if future is None: + raise RuntimeError("No Elastic EP async method is active") + if not future.done(): + raise RuntimeError("Elastic EP async method is not done") + self._async_future = None + future.result() def load_model(self) -> None: - ( - expanded_physical_to_logical, - num_logical_experts, - old_num_physical_experts, - ) = self.receive_expert_mapping() - num_physical_experts = expanded_physical_to_logical.shape[1] - self.worker.parallel_config.eplb_config.num_redundant_experts = ( - num_physical_experts - num_logical_experts - ) self.worker.load_model(load_dummy_weights=True) - self.worker.model_runner.setup_eplb_from_mapping( - expanded_physical_to_logical, old_num_physical_experts - ) - self._set_eplb_suppressed(True) def create_standby_groups( - self, reconfig_request: ReconfigureDistributedRequest + self, reconfig_request: ReconfigureDistributedRequest, use_all2all: bool ) -> None: self.reconfig_request = reconfig_request new_dp_size = reconfig_request.new_data_parallel_size old_dp_size = get_dp_group().world_size - world_size = self.worker.vllm_config.parallel_config.world_size + parallel_config = self.worker.vllm_config.parallel_config + world_size = parallel_config.world_size new_world_size_across_dp = world_size * new_dp_size - updated_config = copy.copy(self.worker.vllm_config) - updated_config.parallel_config = copy.deepcopy( - self.worker.vllm_config.parallel_config + create_standby_groups( + new_dp_size=new_dp_size, + new_world_size_across_dp=new_world_size_across_dp, + master_ip=reconfig_request.new_data_parallel_master_ip, + coord_store_port=reconfig_request.coord_store_port, + use_all2all=use_all2all, + enable_eplb=parallel_config.enable_eplb, ) - updated_config.parallel_config.data_parallel_size = new_dp_size - with set_current_vllm_config(updated_config): - create_standby_groups( - new_dp_size=new_dp_size, - new_world_size_across_dp=new_world_size_across_dp, - master_ip=reconfig_request.new_data_parallel_master_ip, - coord_store_port=reconfig_request.coord_store_port, - enable_eplb=updated_config.parallel_config.enable_eplb, - ) - if new_dp_size > old_dp_size: - self._set_eplb_suppressed(True) - eplb_state = self.worker.model_runner.eplb_state - if eplb_state is not None: - eplb_state.drain_async() - elif new_dp_size < old_dp_size: - self._stage_standby_moe_quant_methods() + if new_dp_size < old_dp_size: + self.stage_standby_moe_quant_methods() def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None: standby_dp_group = get_standby_dp_group() @@ -265,6 +279,7 @@ class ElasticEPScalingExecutor: model_config = self.worker.model_runner.model_config eplb_state = self.worker.model_runner.eplb_state assert eplb_state is not None + eplb_state.drain_async() eplb_model_state = eplb_state.model_states[model_config.compute_hash()] physical_to_logical = eplb_model_state.physical_to_logical_map num_physical_experts = physical_to_logical.shape[1] @@ -278,10 +293,6 @@ class ElasticEPScalingExecutor: src_rank=0, device=self.worker.device, ) - # New workers enter load_model after receiving the expert mapping. - # Stage replacement MoE kernels before returning to the state machine - # so existing ranks can participate in collective EP comm creation. - self._stage_standby_moe_quant_methods() def _make_eep_moe_config(self, module, dp_group, ep_group): parallel_config = self.worker.vllm_config.parallel_config @@ -300,7 +311,7 @@ class ElasticEPScalingExecutor: moe_parallel_config=moe_parallel_config, ) - def _stage_standby_moe_quant_methods(self) -> None: + def stage_standby_moe_quant_methods(self) -> None: standby_dp_group = get_standby_dp_group() standby_ep_group = get_standby_ep_group() model = self.worker.model_runner.get_model() @@ -500,26 +511,6 @@ class ElasticEPScalingExecutor: compilation_counter.stock_torch_compile_count += 1 self.worker.model_runner.model.compile(fullgraph=True, backend=backend) - multi_block_table = self.worker.model_runner.input_batch.block_table - saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] - for bt in multi_block_table.block_tables: - saved_block_tables.append( - (bt.block_table.gpu.clone(), bt.block_table.cpu.clone()) - ) - multi_block_table.clear() - - unlock_workspace() - self.worker.compile_or_warm_up_model() - lock_workspace() - - for bt, (saved_gpu, saved_cpu) in zip( - multi_block_table.block_tables, saved_block_tables - ): - bt.block_table.gpu.copy_(saved_gpu) - bt.block_table.cpu.copy_(saved_cpu) - if new_dp_size < old_dp_size: - self._set_eplb_suppressed(False) - def _perform_eplb_reshuffle( self, rank_mapping: dict[int, int] | None = None ) -> None: @@ -553,12 +544,25 @@ class ElasticEPScalingExecutor: if get_ep_group().rank == 0: logger.info("[Elastic EP] Expert resharding completed") - def perform_eplb_reshuffle(self) -> None: + def commit_scale_up(self, is_existing_worker: bool) -> None: + if is_existing_worker: + self.broadcast_expert_mapping() + self.switch_and_prepare() + else: + mapping, _, num_valid_experts = self.receive_expert_mapping() + self.worker.model_runner.setup_eplb_from_mapping(mapping, num_valid_experts) self._perform_eplb_reshuffle() - self._set_eplb_suppressed(False) + self.warm_and_capture() + + def commit_scale_down(self, new_dp_size: int, removing: bool) -> None: + self.perform_scale_down_eplb_reshuffle(new_dp_size) + if removing: + self.switch_and_remove() + else: + self.switch_and_prepare() + self.warm_and_capture() def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None: - self._set_eplb_suppressed(True) eplb_state = self.worker.model_runner.eplb_state if eplb_state is not None: eplb_state.drain_async() @@ -599,12 +603,17 @@ class ElasticEPScalingExecutor: ) model = self.worker.model_runner.get_model() + expert_weights = [ + module.get_expert_weights() + for module in model.modules() + if is_moe_layer(module) + ] batch_transfer_weights( model=model, is_sender=False, peer_rank=sender_rank, dp_group=dp_group, - expert_weights=model.expert_weights, + expert_weights=expert_weights, ) torch.accelerator.synchronize() @@ -643,14 +652,17 @@ class ElasticEPScalingExecutor: with set_current_vllm_config(self.worker.vllm_config): prepare_communication_buffer_for_model(self.worker.model_runner.get_model()) - def rewarm_workspace(self) -> None: + def warmup_local_kernels(self) -> None: + with set_current_vllm_config(self.worker.vllm_config): + kernel_warmup(self.worker, process_local_only=True) + + def warm_and_capture(self) -> None: # Must run on every DP sibling in lockstep: _dummy_run calls # coordinate_batch_across_dp whenever data_parallel_size > 1 # (gpu_model_runner.py:3663), which deadlocks if any rank skips it. - # Save and clear block tables so profile_run/compile_or_warm_up_model - # don't write dummy slot mappings into real KV-cache blocks (mirrors - # switch_and_prepare's pattern). + # Save and clear block tables so the dummy MoE forward doesn't + # write dummy slot mappings into real KV-cache blocks. multi_block_table = self.worker.model_runner.input_batch.block_table saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = [] for bt in multi_block_table.block_tables: @@ -660,19 +672,16 @@ class ElasticEPScalingExecutor: multi_block_table.clear() # _ensure_workspace_size allocates a fresh tensor on grow, leaving - # captured CUDA graphs with stale data pointers; drop graphs before - # re-warm so captures realign with the resized buffer. + # any captured CUDA graph with a stale data pointer; drop graphs + # before re-warm so captures realign with the resized buffer. self._release_cuda_graphs() unlock_workspace() - # Grow the MoE workspace at max_num_tokens. - # compile_or_warm_up_model alone only exercises cudagraph-capture - # sizes (≤64 tokens for this test) and leaves the workspace at - # ~10-14 MB; the post-all-to-all per-rank token count under real - # post-reshuffle routing needs hundreds of MB. Use _dummy_run - # directly (rather than profile_run) with skip_eplb=True so dummy - # routing doesn't pollute the just-rebalanced EPLB stats — same - # convention compile_or_warm_up_model itself uses. + # Grow the MoE workspace at max_num_tokens. compile_or_warm_up_model + # alone only exercises cudagraph-capture sizes and can leave the + # workspace too small for post-reshuffle routing. Use _dummy_run + # directly with skip_eplb=True so dummy routing doesn't pollute the + # just-rebalanced EPLB stats. runner = self.worker.model_runner runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) self.worker.compile_or_warm_up_model() diff --git a/vllm/distributed/elastic_ep/elastic_state.py b/vllm/distributed/elastic_ep/elastic_state.py index 256efe46a4a..f33fd90e863 100644 --- a/vllm/distributed/elastic_ep/elastic_state.py +++ b/vllm/distributed/elastic_ep/elastic_state.py @@ -1,18 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import enum -import time import weakref -from datetime import timedelta -from typing import TYPE_CHECKING, Literal, TypeAlias +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, Literal, TypeAlias import torch.distributed from vllm.config import ParallelConfig from vllm.distributed import ( - sched_yield, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.utils import get_cached_tcp_store_client from vllm.logger import init_logger from vllm.v1.engine import ( EEPNotificationType, @@ -31,35 +30,29 @@ WorkerType = Literal["existing", "new", "removing"] class ScaleUpExistingEngineState(enum.IntEnum): - WAIT_NEW_CORE_ENGINES_INIT = 0 - CREATE_STANDBY_GROUPS = 1 - TRANSFER_EXPERT_MAPPING = 2 - WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3 - TRANSFER_WEIGHTS = 4 - SYNC_KV_CACHE_MEMORY_SIZE = 5 - SWITCH_AND_PREPARE = 6 - EPLB_RESHUFFLE = 7 - COMPLETE = 8 + CREATE_STANDBY_GROUPS = 0 + STAGE_QUANT_METHODS = 1 + TRANSFER_WEIGHTS = 2 + SYNC_KV_CACHE_MEMORY_SIZE = 3 + COMMIT_SCALE_UP = 4 # Blocks forward passes. + COMPLETE = 5 class ScaleUpNewEngineState(enum.IntEnum): PRE_KV_INIT = 0 PREPARE = 1 - EPLB_RESHUFFLE = 2 - COMPLETE = 3 + COMPLETE = 2 class ScaleDownRemainingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - SWITCH_AND_PREPARE = 2 - COMPLETE = 3 + COMMIT_SCALE_DOWN = 1 # Blocks forward passes. + COMPLETE = 2 class ScaleDownRemovingEngineState(enum.IntEnum): PREPARE = 0 - EPLB_RESHUFFLE = 1 - COMPLETE = 2 + COMPLETE = 1 EngineState: TypeAlias = ( @@ -70,15 +63,6 @@ EngineState: TypeAlias = ( ) -class _BarrierTimeoutError(RuntimeError): - """ - Exception raised for timeout - in the first stage of our two-staged - TCPStore based barrier to synchronize the - execution of all engines in the DP group. - """ - - class ElasticEPScalingState: def __init__( self, @@ -94,20 +78,24 @@ class ElasticEPScalingState: self.engine_core_ref = weakref.ref(engine_core) self.vllm_config = vllm_config self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None - self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None self.new_parallel_config: ParallelConfig = new_parallel_config self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None self.worker_type = worker_type self.scale_type = scale_type self.reconfig_request = reconfig_request - + self.commit_requested = False + self._prepare_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="ElasticEPPrepare" + ) + self._prepare_future: Future[Any] | None = None + self._new_dp_sync: tuple[object, Any] | None = None self.state: EngineState if scale_type == "scale_up": self.state = ( ScaleUpNewEngineState.PRE_KV_INIT if worker_type == "new" - else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT + else ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS ) else: self.state = ( @@ -130,6 +118,31 @@ class ElasticEPScalingState: raise RuntimeError("Engine core has been garbage collected") return engine_core + def _collective_rpc(self, *args, **kwargs): + return self.model_executor.collective_rpc(*args, **kwargs) + + def _execute_async(self, execute_method: str, *args) -> bool: + if self._prepare_future is None: + done_keys = self._collective_rpc( + "elastic_ep_execute", + args=("start_async", execute_method, *args), + ) + assert self.reconfig_request is not None + coord_store = get_cached_tcp_store_client( + self.reconfig_request.new_data_parallel_master_ip, + self.reconfig_request.coord_store_port, + ) + self._prepare_future = self._prepare_executor.submit( + coord_store.wait, done_keys + ) + if not self._prepare_future.done(): + return False + + self._prepare_future.result() + self._collective_rpc("elastic_ep_execute", args=("clear_async",)) + self._prepare_future = None + return True + def progress(self) -> bool: if self.scale_type == "scale_up": return ( @@ -149,157 +162,43 @@ class ElasticEPScalingState: assert self.progress() assert self.state == ScaleUpNewEngineState.PREPARE - def _execute_tcp_store_barrier( - self, dp_store, group_rank, group_size, barrier_id, timeout=None - ): - arrival_key = f"arrival_{barrier_id}_{group_rank}" - dp_store.set(arrival_key, b"1") - - start_time = time.time() - processes_arrived: set[int] = set() - - while len(processes_arrived) < group_size: - if ( - timeout is not None - and time.time() - start_time > timeout.total_seconds() - ): - raise _BarrierTimeoutError( - f"Barrier timed out after {timeout.total_seconds()} seconds" - ) - - for i in range(group_size): - if i in processes_arrived: - continue - - key = f"arrival_{barrier_id}_{i}" - present = dp_store.check([key]) - if present: - processes_arrived.add(i) - - if len(processes_arrived) < group_size: - sched_yield() - - def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool: - """ - Execute a two-staged barrier to synchronize all engines in the DP group. - - Some DP EngineCores may receive the reconfiguration notifications - later than others, and already proceed to engine step (model forward) - in the busy loop. - In this case, EngineCores that already proceed to reconfiguration - should skip reconfiguration and execute model forward for one more - step, so in the next step, all EngineCores will be synchronized. - We use a two-staged barrier to achieve this. The first time each - EngineCore executes the barrier, if a timeout is reached before the - barrier completes, that means some EngineCores have already entered - engine step. The EngineCores that timed out will then proceed to - engine step, and will synchronize with the other EngineCores in the - next step with a barrier without timeout. - """ - dp_group = self.new_dp_group if use_new_group else self.old_dp_group - dp_store = self.new_dp_store if use_new_group else self.old_dp_store - assert dp_group is not None and dp_store is not None - - group_rank = dp_group.rank() - group_size = dp_group.size() - barrier_id = f"eep_barrier_{barrier_name}" - sync_key = f"{barrier_id}_sync" - - # TODO(yongji): figure out appropriate timeout for the barrier - timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5) - - try: - self._execute_tcp_store_barrier( - dp_store, group_rank, group_size, barrier_id, timeout=timeout - ) - torch.distributed.barrier(dp_group) - if group_rank == 0: - dp_store.delete_key(sync_key) - for i in range(group_size): - dp_store.delete_key(f"arrival_{barrier_id}_{i}") - return True - except _BarrierTimeoutError as e: - if timeout is None: - raise RuntimeError("Unexpected timeout encountered") from e - dp_store.compare_set(sync_key, "", b"1") - return False - def _progress_existing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None - if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT: - return False - - elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: - # NOTE(yongji): wait for all existing workers to receive the request - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + if state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS: + if not self._create_standby_groups(): return False - if not self._staged_barrier( - use_new_group=False, barrier_name="create_standby_groups" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._create_standby_groups() - self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING + self.state = ScaleUpExistingEngineState.STAGE_QUANT_METHODS return True - elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING: - self._transfer_expert_mapping() - self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT + elif state == ScaleUpExistingEngineState.STAGE_QUANT_METHODS: + if not self._execute_async("stage_standby_moe_quant_methods"): + return False + self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS return True - elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT: - return False - elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + if not self._transfer_weights(): return False - if not self._staged_barrier( - use_new_group=False, barrier_name="transfer_weights" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._transfer_weights() self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE return True elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE: - self._sync_kv_cache_memory_size() - self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE + if not self._sync_kv_cache_memory_size(): + return False + self.state = ScaleUpExistingEngineState.COMMIT_SCALE_UP + self._mark_ready_for_switch() return True - elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE: - self._switch_and_prepare() - self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE - assert self.new_dp_store is not None - self.new_dp_store.add("eep_barrier_engine_count", 1) - return True - - elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE: - assert self.new_dp_group is not None and self.new_dp_store is not None - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): + elif state == ScaleUpExistingEngineState.COMMIT_SCALE_UP: + if not self.commit_requested: return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): - return False - if self.new_dp_group.rank() == 0: - self.new_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle() + self._commit_new_dp_group() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", True)) self.state = ScaleUpExistingEngineState.COMPLETE self._update_parallel_config() + self._send_reconfigure_finished() return True else: @@ -311,22 +210,17 @@ class ElasticEPScalingState: assert self.new_dp_group is not None and self.new_dp_store is not None if state == ScaleUpNewEngineState.PRE_KV_INIT: - self.engine_core._eep_send_engine_core_notification( - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("receive_weights",) - ) + self._collective_rpc("elastic_ep_execute", args=("receive_weights",)) self.engine_core.available_gpu_memory_for_kv_cache = ( ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1) ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("prepare_new_worker",) - ) + self._collective_rpc("elastic_ep_execute", args=("prepare_new_worker",)) self.state = ScaleUpNewEngineState.PREPARE return True elif state == ScaleUpNewEngineState.PREPARE: + self._collective_rpc("elastic_ep_execute", args=("warmup_local_kernels",)) + self._mark_ready_for_switch() tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu") torch.distributed.all_reduce( tensor, @@ -337,22 +231,7 @@ class ElasticEPScalingState: self.engine_core.engines_running = bool(data[0]) self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) - self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE - self.new_dp_store.add("eep_barrier_engine_count", 1) - return True - - elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE: - if ( - int(self.new_dp_store.get("eep_barrier_engine_count")) - < self.new_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=True, barrier_name="eplb_reshuffle" - ): - return False - assert self.new_dp_group.rank() > 0 - self._eplb_reshuffle() + self._collective_rpc("elastic_ep_execute", args=("commit_scale_up", False)) self.state = ScaleUpNewEngineState.COMPLETE return True @@ -362,38 +241,23 @@ class ElasticEPScalingState: def _progress_remaining_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemainingEngineState.PREPARE: - self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True + if self._create_standby_groups(): + self.state = ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + self._mark_ready_for_switch() + return True + return False - elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): + elif state == ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN: + if not self.commit_requested: return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): - return False - if self.old_dp_group.rank() == 0: - self.old_dp_store.delete_key("eep_barrier_engine_count") - self._eplb_reshuffle_before_scale_down() - self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE - # NOTE(yongji): currently, after EPLB reshuffle - # that redistributes experts to remaining workers, workers - # to be removed will immediately initiate shutdown; - # existing workers can no longer execute forward steps using - # the old setup. In the future, we may keep - # the removing workers alive a bit longer, - # e.g., to drain in-batch requests. - self._create_standby_groups() - self._switch_and_prepare() + self._commit_scale_down(removing=False) + self._commit_new_dp_group() self._update_parallel_config() self.state = ScaleDownRemainingEngineState.COMPLETE + self._send_reconfigure_finished() return True else: @@ -402,26 +266,11 @@ class ElasticEPScalingState: def _progress_removing_engine(self) -> bool: state = self.state - assert self.old_dp_group is not None and self.old_dp_store is not None + assert self.old_dp_group is not None if state == ScaleDownRemovingEngineState.PREPARE: - self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE - self.old_dp_store.add("eep_barrier_engine_count", 1) - return True - - if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE: - if ( - int(self.old_dp_store.get("eep_barrier_engine_count")) - < self.old_dp_group.size() - ): - return False - if not self._staged_barrier( - use_new_group=False, barrier_name="eplb_reshuffle" - ): - return False assert self.old_dp_group.rank() > 0 - self._eplb_reshuffle_before_scale_down() - self._switch_and_remove() + self._commit_scale_down(removing=True) self.state = ScaleDownRemovingEngineState.COMPLETE self.engine_core._eep_send_engine_core_notification( EEPNotificationType.SHUTDOWN_COMPLETE @@ -432,22 +281,22 @@ class ElasticEPScalingState: assert self.state == ScaleDownRemovingEngineState.COMPLETE return True - def handle_notification(self, notification_type: EEPNotificationType): - assert self.worker_type != "new" - assert self.old_dp_store is not None - if ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY - and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS - elif ( - notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY - and self.state - == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT - ): - self.old_dp_store.add("eep_barrier_engine_count", 1) - self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS + def is_ready_for_switch(self) -> bool: + return self.worker_type == "existing" and ( + self.state is ScaleUpExistingEngineState.COMMIT_SCALE_UP + or self.state is ScaleDownRemainingEngineState.COMMIT_SCALE_DOWN + ) + + @property + def ready_key(self) -> str: + return f"eep_ready/{self.engine_core.dp_rank}" + + def _mark_ready_for_switch(self) -> None: + parallel_config = self.new_parallel_config + get_cached_tcp_store_client( + parallel_config.data_parallel_master_ip, + parallel_config._coord_store_port, + ).set(self.ready_key, b"1") def is_complete(self) -> bool: if self.scale_type == "scale_up": @@ -462,50 +311,78 @@ class ElasticEPScalingState: else self.state == ScaleDownRemainingEngineState.COMPLETE ) - def _create_standby_groups(self): + def _init_new_dp_group(self) -> tuple[Any, Any]: + return self.new_parallel_config.stateless_init_dp_group(return_store=True) + + def _ensure_new_dp_group(self) -> bool: + if self.new_dp_group is not None: + return True + + if self._prepare_future is None: + self._prepare_future = self._prepare_executor.submit( + self._init_new_dp_group + ) + if not self._prepare_future.done(): + return False + + self.new_dp_group, self.new_dp_store = self._prepare_future.result() + self._prepare_future = None + return True + + def _create_standby_groups(self) -> bool: assert self.old_dp_group is not None - self.new_dp_group, self.new_dp_store = ( - self.new_parallel_config.stateless_init_dp_group(return_store=True) - ) - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request) - ) + if not self._ensure_new_dp_group(): + return False + if not self._execute_async( + "create_standby_groups", + self.reconfig_request, + self.new_parallel_config.use_all2all, + ): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Created standby communication groups") + return True - def _transfer_weights(self): + def _transfer_weights(self) -> bool: assert self.reconfig_request is not None and self.old_dp_group is not None old_dp_size = self.old_dp_group.size() new_dp_size = self.reconfig_request.new_data_parallel_size - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size) - ) + if not self._execute_async("transfer_weights", old_dp_size, new_dp_size): + return False if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Transferred weights to new workers") + return True - def _transfer_expert_mapping(self): - assert self.old_dp_group is not None - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("broadcast_expert_mapping",) - ) - if self.old_dp_group.rank() == 0: - logger.info("[Elastic EP] Broadcasted expert mapping to new workers") - - def _sync_kv_cache_memory_size(self): + def _sync_kv_cache_memory_size(self) -> bool: assert self.engine_core.available_gpu_memory_for_kv_cache > 0 assert self.new_dp_group is not None and self.old_dp_group is not None - ParallelConfig.sync_kv_cache_memory_size( - self.new_dp_group, - self.engine_core.available_gpu_memory_for_kv_cache, - ) + + if self._new_dp_sync is None: + tensor = torch.tensor( + [self.engine_core.available_gpu_memory_for_kv_cache], + dtype=torch.int64, + device="cpu", + ) + work = torch.distributed.all_reduce( + tensor, + op=torch.distributed.ReduceOp.MIN, + group=self.new_dp_group, + async_op=True, + ) + self._new_dp_sync = (tensor, work) + return False + + _, work = self._new_dp_sync + if not work.is_completed(): + return False + work.wait() + self._new_dp_sync = None if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] Synced KV cache memory size to new workers") + return True - def _switch_and_prepare(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_prepare",) - ) + def _commit_new_dp_group(self): old_dp_group = self.old_dp_group stateless_destroy_torch_distributed_process_group(old_dp_group) assert self.new_dp_group is not None @@ -529,41 +406,28 @@ class ElasticEPScalingState: self.engine_core.current_wave = int(data[1]) self.engine_core.step_counter = int(data[2]) if new_dp_group.rank() == 0: + logger.info("[Elastic EP] Switched to new setup") + + def _send_reconfigure_finished(self): + assert self.new_dp_group is not None + if self.new_dp_group.rank() == 0: self.engine_core._eep_send_engine_core_notification( EEPNotificationType.RECONFIGURE_FINISHED ) - logger.info("[Elastic EP] Switched to new setup") - def _eplb_reshuffle(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("perform_eplb_reshuffle",) - ) - # Reshuffle changes per-rank token routing; the locked MoE workspace - # may now be too small. Rewarm covers both new and existing engines. - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("rewarm_workspace",) - ) - assert self.new_dp_group is not None - if self.new_dp_group.rank() == 0: - logger.info("[Elastic EP] EPLB reshuffle completed") - - def _eplb_reshuffle_before_scale_down(self): + def _commit_scale_down(self, removing: bool): assert self.reconfig_request is not None and self.old_dp_group is not None - self.model_executor.collective_rpc( + self._collective_rpc( "elastic_ep_execute", args=( - "perform_scale_down_eplb_reshuffle", + "commit_scale_down", self.reconfig_request.new_data_parallel_size, + removing, ), ) if self.old_dp_group.rank() == 0: logger.info("[Elastic EP] EPLB reshuffle completed") - def _switch_and_remove(self): - self.model_executor.collective_rpc( - "elastic_ep_execute", args=("switch_and_remove",) - ) - def _update_parallel_config(self): assert self.reconfig_request is not None reconfig_request = self.reconfig_request diff --git a/vllm/distributed/elastic_ep/standby_state.py b/vllm/distributed/elastic_ep/standby_state.py index 846793a955f..1892f3e7942 100644 --- a/vllm/distributed/elastic_ep/standby_state.py +++ b/vllm/distributed/elastic_ep/standby_state.py @@ -39,6 +39,7 @@ def create_standby_groups( new_world_size_across_dp: int, master_ip: str, coord_store_port: int, + use_all2all: bool, enable_eplb: bool = True, backend: str | None = None, ) -> None: @@ -86,7 +87,7 @@ def create_standby_groups( ) standby_ep_ranks = [x.tolist() for x in standby_ep_ranks] _STANDBY_EP = _init_stateless_group( - standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store + standby_ep_ranks, "ep", master_ip, backend, coord_store, use_all2all=use_all2all ) if enable_eplb: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 4284a609d67..a90e8acbcad 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -414,6 +414,7 @@ class GroupCoordinator: use_device_communicator: bool, # whether to use device communicator use_message_queue_broadcaster: bool = False, group_name: str | None = None, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -508,6 +509,7 @@ class GroupCoordinator: device=self.device, device_group=self.device_group, unique_name=self.unique_name, + use_all2all=use_all2all, ) from vllm.distributed.device_communicators.shm_broadcast import MessageQueue @@ -1321,6 +1323,7 @@ def init_model_parallel_group( use_message_queue_broadcaster: bool = False, group_name: str | None = None, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> GroupCoordinator: return GroupCoordinator( group_ranks=group_ranks, @@ -1329,6 +1332,7 @@ def init_model_parallel_group( use_device_communicator=use_device_communicator, use_message_queue_broadcaster=use_message_queue_broadcaster, group_name=group_name, + use_all2all=use_all2all, ) @@ -1339,6 +1343,7 @@ def _init_stateless_group( backend: str, coord_store: Store, use_device_communicator: bool = True, + use_all2all: bool = False, ) -> "StatelessGroupCoordinator": """Create a StatelessGroupCoordinator with the given parameters.""" from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator @@ -1354,6 +1359,7 @@ def _init_stateless_group( coord_store=coord_store, global_rank=world.rank, global_world_size=world.world_size, + use_all2all=use_all2all, ) @@ -1924,6 +1930,7 @@ def initialize_model_parallel( .unbind(0) ) group_ranks = [x.tolist() for x in group_ranks] + use_all2all = parallel_config.use_all2all if enable_elastic_ep: _EP = _init_stateless_group( group_ranks, @@ -1931,10 +1938,15 @@ def initialize_model_parallel( parallel_config.data_parallel_master_ip, backend, coord_store=coord_store, + use_all2all=use_all2all, ) else: _EP = init_model_parallel_group( - group_ranks, get_world_group().local_rank, backend, group_name="ep" + group_ranks, + get_world_group().local_rank, + backend, + group_name="ep", + use_all2all=use_all2all, ) # Create EPLB group with the same ranks as EP if EPLB is enabled. diff --git a/vllm/distributed/stateless_coordinator.py b/vllm/distributed/stateless_coordinator.py index 5f4597d07cb..38c74a97c55 100644 --- a/vllm/distributed/stateless_coordinator.py +++ b/vllm/distributed/stateless_coordinator.py @@ -79,6 +79,7 @@ class StatelessGroupCoordinator(GroupCoordinator): host: str = "127.0.0.1", global_rank: int = 0, global_world_size: int = 1, + use_all2all: bool = False, ): group_name = group_name or "anonymous" self.unique_name = _get_unique_name(group_name) @@ -191,6 +192,7 @@ class StatelessGroupCoordinator(GroupCoordinator): global_ranks=self.ranks, global_world_size=global_world_size, tcp_store_group=self.tcp_store_group, + use_all2all=use_all2all, ) self.mq_broadcaster = None diff --git a/vllm/entrypoints/serve/elastic_ep/api_router.py b/vllm/entrypoints/serve/elastic_ep/api_router.py index e711a257ddd..02a24250905 100644 --- a/vllm/entrypoints/serve/elastic_ep/api_router.py +++ b/vllm/entrypoints/serve/elastic_ep/api_router.py @@ -12,10 +12,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, ) -from vllm.entrypoints.serve.elastic_ep.middleware import ( - get_scaling_elastic_ep, - set_scaling_elastic_ep, -) +from vllm.entrypoints.serve.elastic_ep.middleware import get_scaling_elastic_ep from vllm.entrypoints.serve.utils.api_utils import validate_json_request from vllm.logger import init_logger @@ -64,8 +61,6 @@ async def scale_elastic_ep(raw_request: Request): status_code=400, detail="drain_timeout must be a positive integer" ) - # Set scaling flag to prevent new requests - set_scaling_elastic_ep(True) client = engine_client(raw_request) try: await client.scale_elastic_ep(new_data_parallel_size, drain_timeout) @@ -83,8 +78,6 @@ async def scale_elastic_ep(raw_request: Request): except Exception as e: logger.error("Scale failed: %s", e) raise HTTPException(status_code=500, detail="Scale failed") from e - finally: - set_scaling_elastic_ep(False) @router.post("/is_scaling_elastic_ep") diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e461dae0bb8..b2c989b295c 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -93,7 +93,7 @@ def _warmup_ll_bf16_router_gemm(model: torch.nn.Module) -> None: ) -def kernel_warmup(worker: "Worker"): +def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( minimax_m3_msa_warmup, ) @@ -118,6 +118,22 @@ def kernel_warmup(worker: "Worker"): ) # Run next so input-prep kernels JIT against pristine runner state. + if worker.vllm_config.kernel_config.enable_jit_warmup: + fa4_cutedsl_warmup(worker) + sparse_mla_triton_warmup(worker) + + if current_platform.has_device_capability(90): + _warmup_ll_bf16_router_gemm(worker.get_model()) + + if worker.vllm_config.kernel_config.enable_cutedsl_warmup: + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 + cutedsl_warmup() + + if process_local_only: + return + flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) @@ -143,9 +159,6 @@ def kernel_warmup(worker: "Worker"): elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) - if current_platform.has_device_capability(90): - _warmup_ll_bf16_router_gemm(worker.get_model()) - # FlashInfer attention warmup # Only warmup if the model has FlashInfer attention groups # and is not a pooling model @@ -178,16 +191,6 @@ def kernel_warmup(worker: "Worker"): create_mixed_batch=True, ) - if worker.vllm_config.kernel_config.enable_cutedsl_warmup: - # TODO(roberto): Remove after registered CuTeDSL warmups are migrated - # to the shared JIT warmup infrastructure. - # https://github.com/vllm-project/vllm/pull/47451 - cutedsl_warmup() - - if worker.vllm_config.kernel_config.enable_jit_warmup: - fa4_cutedsl_warmup(worker) - sparse_mla_triton_warmup(worker) - def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None: diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index e80be0e45d7..83033ecf81d 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -35,8 +35,6 @@ FT_STATUS_CALL_ID = -2 class EEPNotificationType(enum.Enum): - NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY" - NEW_CORE_ENGINES_WEIGHTS_INIT_READY = "NEW_CORE_ENGINES_WEIGHTS_INIT_READY" RECONFIGURE_FINISHED = "RECONFIGURE_FINISHED" SHUTDOWN_COMPLETE = "SHUTDOWN_COMPLETE" diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 5c2e01cf44b..e8d33c961dc 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -109,6 +109,7 @@ class AsyncLLM(EngineClient): maybe_register_config_serialize_by_value() self.vllm_config = vllm_config + self._elastic_ep_lock = asyncio.Lock() self.model_config = vllm_config.model_config self.observability_config = vllm_config.observability_config @@ -998,17 +999,27 @@ class AsyncLLM(EngineClient): "waiting for requests to drain." ) + async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None: + try: + logger.info( + "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, " + "waiting for requests to drain before scaling" + ) + await self.wait_for_requests_to_drain(drain_timeout) + except BaseException: + set_scaling_elastic_ep(False) + raise + async def scale_elastic_ep( self, new_data_parallel_size: int, drain_timeout: int = 300 ): - """ - Scale up or down the data parallel size by adding or removing - engine cores. - Args: - new_data_parallel_size: The new number of data parallel workers - drain_timeout: - Maximum time to wait for requests to drain (seconds) - """ + """Scale the elastic EP data parallel size.""" + async with self._elastic_ep_lock: + await self._scale_elastic_ep(new_data_parallel_size, drain_timeout) + + async def _scale_elastic_ep( + self, new_data_parallel_size: int, drain_timeout: int + ) -> None: old_data_parallel_size = self.vllm_config.parallel_config.data_parallel_size if old_data_parallel_size == new_data_parallel_size: logger.info( @@ -1017,12 +1028,7 @@ class AsyncLLM(EngineClient): ) return - if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: - logger.info( - "VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, " - "waiting for requests to drain before scaling" - ) - await self.wait_for_requests_to_drain(drain_timeout) + await self.engine_core.prepare_elastic_ep(new_data_parallel_size) # recreate stat loggers if new_data_parallel_size > old_data_parallel_size and self.log_stats: @@ -1042,11 +1048,12 @@ class AsyncLLM(EngineClient): self.logger_manager.log_engine_initialized() set_scaling_elastic_ep(True) - try: - await self.engine_core.scale_elastic_ep(new_data_parallel_size) - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - finally: - set_scaling_elastic_ep(False) + if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS: + await self._drain_requests_for_elastic_ep(drain_timeout) + + await self.engine_core.commit_elastic_ep() + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + set_scaling_elastic_ep(False) async def handle_fault( self, fault_tolerance_request: FaultToleranceRequest diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index 2f3b03636d7..d7f05cffc8a 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -376,6 +376,9 @@ class DPCoordinatorProc: eng_index = outputs.engine_index scheduler_stats = outputs.scheduler_stats if scheduler_stats: + # Elastic EP stats may arrive while the engine list changes. + if eng_index >= len(self.engines): + continue # 1. Updated request load stats - update our local # state with these. stats = self.engines[eng_index].request_counts diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 9917f810b5b..ecac92f5fe0 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -327,8 +327,9 @@ class EngineCore: vllm_config.validate_block_size() - # Initialize kv cache and warmup the execution self.model_executor.initialize_from_config(kv_cache_configs) + if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: + self.model_executor.compile_or_warm_up_model() elapsed = time.time() - start compile_time = vllm_config.compilation_config.compilation_time @@ -993,9 +994,7 @@ class EngineCore: raise NotImplementedError def _eep_send_engine_core_notification( - self, - notification_type: EEPNotificationType, - vllm_config: VllmConfig | None = None, + self, notification_type: EEPNotificationType ): raise NotImplementedError @@ -1070,11 +1069,6 @@ class EngineCoreProc(EngineCore): self.addresses = addresses self.process_input_queue_block = True - if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: - self._eep_send_engine_core_notification( - EEPNotificationType.NEW_CORE_ENGINES_INIT_READY, - vllm_config=vllm_config, - ) self._init_data_parallel(vllm_config) super().__init__( @@ -2096,12 +2090,16 @@ class DPEngineCoreProc(EngineCoreProc): self._maybe_publish_request_counts() if self.eep_scaling_state is not None: - _ = self.eep_scaling_state.progress() - if self.eep_scaling_state.is_complete(): - if self.eep_scaling_state.worker_type == "removing": + state = self.eep_scaling_state + if state.commit_requested or not state.is_ready_for_switch(): + state.progress() + if state.is_complete(): + if state.worker_type == "removing": raise SystemExit self.process_input_queue_block = True self.eep_scaling_state = None + elif not state.commit_requested and state.is_ready_for_switch(): + self.process_input_queue_block = True executed = self._process_engine_step() self._maybe_publish_request_counts() @@ -2171,7 +2169,7 @@ class DPEngineCoreProc(EngineCoreProc): def reinitialize_distributed( self, reconfig_request: ReconfigureDistributedRequest - ) -> None: + ) -> str: from copy import deepcopy from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState @@ -2203,7 +2201,10 @@ class DPEngineCoreProc(EngineCoreProc): == ReconfigureRankType.SHUTDOWN_CURRENT_RANK ) - self.eep_scaling_state = ElasticEPScalingState( + if self.eep_scaling_state is not None: + raise RuntimeError("Elastic EP reconfiguration is already active") + + state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2212,30 +2213,34 @@ class DPEngineCoreProc(EngineCoreProc): scale_type="scale_down" if is_scale_down else "scale_up", reconfig_request=reconfig_request, ) + self.eep_scaling_state = state + self.process_input_queue_block = False logger.info( "[Elastic EP] Received reconfiguration request and starting scaling up/down" ) + return state.ready_key + + def commit_prepared_elastic_ep(self) -> None: + state = self.eep_scaling_state + if state is None or state.commit_requested or not state.is_ready_for_switch(): + raise RuntimeError("No prepared Elastic EP reconfiguration is ready") + state.commit_requested = True + self.process_input_queue_block = False + logger.info("[Elastic EP] Committing prepared reconfiguration") def _eep_send_engine_core_notification( - self, - notification_type: EEPNotificationType, - vllm_config: VllmConfig | None = None, + self, notification_type: EEPNotificationType ): """ Send notifications to EngineCoreClient, which can then forward the notifications to other engine core processes. It is used for: - 1) In scale up: new core engines to notify existing core engines - that they are ready; - 2) In scale down: removing core engines to notify EngineCoreClient + 1) In scale down: removing core engines to notify EngineCoreClient so EngineCoreClient can release their ray placement groups; - 3) Both scale up/down: to notify EngineCoreClient that existing + 2) Both scale up/down: to notify EngineCoreClient that existing core engines have already switched to the new parallel setup. """ - if vllm_config is None: - dp_rank = self.vllm_config.parallel_config.data_parallel_rank - else: - dp_rank = vllm_config.parallel_config.data_parallel_rank + dp_rank = self.vllm_config.parallel_config.data_parallel_rank notification_data = (notification_type.value, dp_rank) outputs = EngineCoreOutputs( utility_output=UtilityOutput( @@ -2257,22 +2262,11 @@ class DPEngineCoreProc(EngineCoreProc): ): socket.send_multipart(encoder.encode(outputs)) - def eep_handle_engine_core_notification( - self, notification_type: str | EEPNotificationType - ): - """ - Handle notification received from EngineCoreClient - (forwarded from new core engines). - """ - assert self.eep_scaling_state is not None - if isinstance(notification_type, str): - notification_type = EEPNotificationType(notification_type) - self.eep_scaling_state.handle_notification(notification_type) - def _eep_scale_up_before_kv_init(self): from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState - self.eep_scaling_state = ElasticEPScalingState( + self.ignore_start_dp_wave = True + state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2281,7 +2275,10 @@ class DPEngineCoreProc(EngineCoreProc): scale_type="scale_up", reconfig_request=None, ) - self.eep_scaling_state.run_pre_kv_init_states() + if self.eep_scaling_state is not None: + raise RuntimeError("Elastic EP reconfiguration is already active") + self.eep_scaling_state = state + state.run_pre_kv_init_states() self.process_input_queue_block = False diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index a6c232b2ab7..9460fdf48f7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -225,7 +225,10 @@ class EngineCoreClient(ABC): running state.""" raise NotImplementedError - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None: + async def commit_elastic_ep(self) -> None: + raise NotImplementedError + + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: raise NotImplementedError async def get_output_async(self) -> EngineCoreOutputs: @@ -1490,6 +1493,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) assert len(self.core_engines) > 1 + self._prepared_elastic_ep: tuple[int, int] | None = None self.eng_start_index = ( len(self.core_engines) * self.client_index @@ -1603,31 +1607,16 @@ class DPLBAsyncMPClient(DPAsyncMPClient): if len(cache.pending_notifications[notification_type]) >= abs( cache.num_new_core_engines ): - if notification_type == EEPNotificationType.SHUTDOWN_COMPLETE: - assert isinstance(self.resources.engine_manager, CoreEngineActorManager) - assert cache.num_new_core_engines < 0 - old_dp_size = len(cache.existing_core_engines) - new_dp_size = old_dp_size + cache.num_new_core_engines - self.resources.engine_manager.scale_down_elastic_ep( - old_dp_size, new_dp_size - ) - else: - await asyncio.gather( - *[ - self._call_utility_async( - "eep_handle_engine_core_notification", - notification_type, - engine=engine, - ) - for engine in cache.existing_core_engines - ] - ) - cache.pending_notifications[notification_type] = set() - if notification_type in [ - EEPNotificationType.SHUTDOWN_COMPLETE, - EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY, - ]: - self.eep_scaling_cache = None + engine_manager = self.resources.engine_manager + assert isinstance(engine_manager, CoreEngineActorManager) + assert cache.num_new_core_engines < 0 + old_dp_size = len(cache.existing_core_engines) + new_dp_size = old_dp_size + cache.num_new_core_engines + engine_manager.scale_down_elastic_ep(old_dp_size, new_dp_size) + self.vllm_config.parallel_config.data_parallel_size_local = len( + engine_manager.local_engine_actors + ) + self.eep_scaling_cache = None async def abort_requests_async(self, request_ids: list[str]) -> None: if not request_ids or self.resources.engine_dead: @@ -1651,31 +1640,46 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) -> None: await self._send_input(EngineCoreRequestType.ABORT, request_ids, engine) - async def scale_elastic_ep(self, new_data_parallel_size: int) -> None: - """Scale elastic EP data parallel size""" + async def commit_elastic_ep(self) -> None: + """Commit prepared elastic EP scaling.""" + prepared = self._prepared_elastic_ep + if prepared is None: + raise RuntimeError("Elastic EP scaling has not been prepared") + new_data_parallel_size, num_redundant_experts = prepared cur_data_parallel_size = len(self.core_engines) - - assert new_data_parallel_size != cur_data_parallel_size, ( - f"new_data_parallel_size {new_data_parallel_size} must be " - f"different from cur_data_parallel_size {cur_data_parallel_size}" + if new_data_parallel_size > cur_data_parallel_size: + await self._commit_scale_up_elastic_ep(new_data_parallel_size) + else: + await self._commit_scale_down_elastic_ep(new_data_parallel_size) + self.vllm_config.parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts ) + self._prepared_elastic_ep = None + async def prepare_elastic_ep(self, new_data_parallel_size: int) -> None: + """Prepare elastic EP scaling without routing requests to new engines.""" + if (prepared := self._prepared_elastic_ep) is not None: + if prepared[0] == new_data_parallel_size: + return + raise RuntimeError("Elastic EP scaling is already prepared") + cur_data_parallel_size = len(self.core_engines) assert self.vllm_config.parallel_config.data_parallel_backend == "ray", ( "Only ray DP backend supports scaling elastic EP" ) - - scale_up = new_data_parallel_size > cur_data_parallel_size - - if scale_up: - await self._scale_up_elastic_ep( - cur_data_parallel_size, new_data_parallel_size - ) + parallel_config = self.vllm_config.parallel_config + num_experts = self.vllm_config.model_config.get_num_experts() + num_redundant_experts = ( + num_experts + parallel_config.eplb_config.num_redundant_experts + ) * new_data_parallel_size // cur_data_parallel_size - num_experts + if new_data_parallel_size < cur_data_parallel_size: + await self._prepare_scale_down_elastic_ep(new_data_parallel_size) else: - await self._scale_down_elastic_ep( - cur_data_parallel_size, new_data_parallel_size + await self._prepare_scale_up_elastic_ep( + new_data_parallel_size, num_redundant_experts ) + self._prepared_elastic_ep = new_data_parallel_size, num_redundant_experts - async def _eep_wait_for_setup_switch_complete(self) -> None: + def _eep_wait_for_setup_switch_complete(self) -> asyncio.Future: """ Wait for core engines to switch to the new setup. @@ -1687,9 +1691,26 @@ class DPLBAsyncMPClient(DPAsyncMPClient): future = asyncio.get_running_loop().create_future() self.utility_results[EEP_NOTIFICATION_CALL_ID] = future self._ensure_output_queue_task() - await future + return future - def _setup_elastic_ep_reconfig_bootstrap(self) -> tuple[str, int]: + def _wait_for_new_engine_ready(self, new_core_engines: list[bytes]) -> None: + new_engine_identities = set(new_core_engines) + sync_input_socket = zmq.Socket.shadow(self.input_socket) + while new_engine_identities: + if not sync_input_socket.poll(timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000): + raise TimeoutError( + f"Timed out waiting for new engine core processes to " + f"start. Waited " + f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by " + f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the " + f"timeout, set the environment variable: " + f"VLLM_ENGINE_READY_TIMEOUT_S=" + ) + identity, payload = sync_input_socket.recv_multipart() + new_engine_identities.discard(identity) + self._apply_ready_response(payload) + + def _setup_elastic_ep_reconfig_bootstrap(self) -> None: from vllm.distributed.utils import create_tcp_store from vllm.utils.network_utils import get_open_ports_list @@ -1709,36 +1730,36 @@ class DPLBAsyncMPClient(DPAsyncMPClient): ) parallel_config._coord_store_port = store.port self._coord_store = store - return ip, store.port - async def _scale_up_elastic_ep( - self, cur_data_parallel_size: int, new_data_parallel_size: int - ) -> None: - """Scale up the data parallel size by creating new engine cores - and reconfiguring existing ones.""" - cur_data_parallel_size = len(self.core_engines) - - self.eep_scaling_cache = ElasticScalingCache( - existing_core_engines=self.core_engines.copy(), - num_new_core_engines=new_data_parallel_size - cur_data_parallel_size, - pending_notifications=dict(), + def _make_reconfig_request( + self, + new_data_parallel_size: int, + rank_type: ReconfigureRankType = ReconfigureRankType.KEEP_CURRENT_RANK, + ) -> ReconfigureDistributedRequest: + parallel_config = self.vllm_config.parallel_config + return ReconfigureDistributedRequest( + new_data_parallel_size=new_data_parallel_size, + new_data_parallel_rank=rank_type, + new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, + new_data_parallel_master_ip=parallel_config.data_parallel_master_ip, + new_data_parallel_master_port=parallel_config.data_parallel_master_port, + new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, + coord_store_port=parallel_config._coord_store_port, ) - parallel_config = self.vllm_config.parallel_config - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap() + async def _prepare_scale_up_elastic_ep( + self, + new_data_parallel_size: int, + num_redundant_experts: int, + ) -> None: + """Prepare scale up by creating new engine cores and reconfiguring + existing ones.""" + self._setup_elastic_ep_reconfig_bootstrap() # Phase 1: Send reconfig messages to existing engines reconfig_futures = [] for engine in self.core_engines: - reconfig_request = ReconfigureDistributedRequest( - new_data_parallel_size=new_data_parallel_size, - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_master_ip=ip, - new_data_parallel_master_port=parallel_config.data_parallel_master_port, - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, - coord_store_port=coord_store_port, - ) + reconfig_request = self._make_reconfig_request(new_data_parallel_size) coro = self._call_utility_async( "reinitialize_distributed", reconfig_request, engine=engine ) @@ -1746,51 +1767,54 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # Phase 2: Create new engines assert isinstance(self.resources.engine_manager, CoreEngineActorManager) - parallel_config.eplb_config.num_redundant_experts = 0 start_new_worker_future = asyncio.to_thread( self.resources.engine_manager.scale_up_elastic_ep, self.vllm_config, new_data_parallel_size, + num_redundant_experts, ) - wait_future = self._eep_wait_for_setup_switch_complete() # Phase 3: Wait for new engines to be created # and reconfig messages to be received await asyncio.gather(start_new_worker_future, *reconfig_futures) + ready_keys = [future.result() for future in reconfig_futures] + ready_keys.extend( + f"eep_ready/{rank}" + for rank in range(len(self.core_engines), new_data_parallel_size) + ) + await asyncio.to_thread(self._coord_store.wait, ready_keys) logger.info("[Elastic EP] Successfully started new engines") - # Create new CoreEngine objects for the new engines - new_engine_identities = set() - for i in range(cur_data_parallel_size, new_data_parallel_size): - new_engine = i.to_bytes(2, "little") - self.core_engines.append(new_engine) - # NOTE(yongji): we don't update lb_engines here, - # we let run_engine_stats_update_task to update it. - new_engine_identities.add(new_engine) + async def _commit_scale_up_elastic_ep(self, new_data_parallel_size: int) -> None: + new_core_engines = [ + rank.to_bytes(2, "little") + for rank in range(len(self.core_engines), new_data_parallel_size) + ] - # Wait for ready messages from new engines on the input socket - sync_input_socket = zmq.Socket.shadow(self.input_socket) - while new_engine_identities: - if not sync_input_socket.poll( - timeout=VLLM_ENGINE_READY_TIMEOUT_S * 1000 # convert to ms - ): - raise TimeoutError( - f"Timed out waiting for new engine core processes to " - f"start. Waited " - f"{VLLM_ENGINE_READY_TIMEOUT_S}s (configured by " - f"VLLM_ENGINE_READY_TIMEOUT_S). To increase the " - f"timeout, set the environment variable: " - f"VLLM_ENGINE_READY_TIMEOUT_S=" - ) - identity, payload = sync_input_socket.recv_multipart() - new_engine_identities.discard(identity) - self._apply_ready_response(payload) + await self.pause_scheduler_async(mode="keep", clear_cache=False) + wait_future = self._eep_wait_for_setup_switch_complete() + finish_futures = [ + asyncio.create_task( + self._call_utility_async("commit_prepared_elastic_ep", engine=engine) + ) + for engine in self.core_engines + ] + try: + await asyncio.gather(*finish_futures) + await wait_future + self._wait_for_new_engine_ready(new_core_engines) + except Exception: + wait_future.cancel() + raise - # NOTE(yongji): Before we schedule any requests on the new workers, - # we should wait for them to switch to the new setup. - await wait_future + self.core_engines.extend(new_core_engines) # Update the parallel config - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + parallel_config = self.vllm_config.parallel_config + parallel_config.data_parallel_size = new_data_parallel_size + if isinstance(self.resources.engine_manager, CoreEngineActorManager): + parallel_config.data_parallel_size_local = len( + self.resources.engine_manager.local_engine_actors + ) # Notify coordinator about scale up through existing # stats_update_task connection self._ensure_stats_update_task() @@ -1803,10 +1827,23 @@ class DPLBAsyncMPClient(DPAsyncMPClient): "[Elastic EP] Scale up completed, new data parallel size: %s", new_data_parallel_size, ) + await self.resume_scheduler_async() - async def _scale_down_elastic_ep( - self, cur_data_parallel_size: int, new_data_parallel_size: int - ) -> None: + async def _prepare_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: + self._setup_elastic_ep_reconfig_bootstrap() + + reconfig_futures = [] + for engine in self.core_engines[:new_data_parallel_size]: + reconfig_request = self._make_reconfig_request(new_data_parallel_size) + coro = self._call_utility_async( + "reinitialize_distributed", reconfig_request, engine=engine + ) + reconfig_futures.append(asyncio.create_task(coro)) + + ready_keys = await asyncio.gather(*reconfig_futures) + await asyncio.to_thread(self._coord_store.wait, ready_keys) + + async def _commit_scale_down_elastic_ep(self, new_data_parallel_size: int) -> None: """Scale down the data parallel size by shutting down and reconfiguring existing engine cores.""" cur_data_parallel_size = len(self.core_engines) @@ -1817,50 +1854,51 @@ class DPLBAsyncMPClient(DPAsyncMPClient): pending_notifications=dict(), ) - parallel_config = self.vllm_config.parallel_config - ip, coord_store_port = self._setup_elastic_ep_reconfig_bootstrap() - + old_core_engines = self.core_engines + # NOTE(yongji): Immediately stop sending requests to the removing engines. + self.core_engines = old_core_engines[:new_data_parallel_size] + self.lb_engines = self.lb_engines[:new_data_parallel_size] removed_dp_size = cur_data_parallel_size - new_data_parallel_size + pause_modes = ["keep"] * new_data_parallel_size + ["abort"] * removed_dp_size + pause_futures = [ + self._call_utility_async("pause_scheduler", mode, False, engine=engine) + for mode, engine in zip(pause_modes, old_core_engines) + ] + await asyncio.gather(*pause_futures) assert isinstance(self.resources.engine_manager, CoreEngineActorManager) self.resources.engine_manager.remove_run_refs_for_scale_down(removed_dp_size) + wait_future = self._eep_wait_for_setup_switch_complete() reconfig_futures = [] - for cur_dp_rank, engine in enumerate(self.core_engines): - reconfig_request = ReconfigureDistributedRequest( - new_data_parallel_size=new_data_parallel_size, - new_data_parallel_rank=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_rank_local=ReconfigureRankType.KEEP_CURRENT_RANK, - new_data_parallel_master_ip=ip, - new_data_parallel_master_port=parallel_config.data_parallel_master_port, - new_data_parallel_master_port_list=parallel_config._data_parallel_master_port_list, - coord_store_port=coord_store_port, - ) - if cur_dp_rank >= new_data_parallel_size: - reconfig_request.new_data_parallel_rank = ( - ReconfigureRankType.SHUTDOWN_CURRENT_RANK + for cur_dp_rank, engine in enumerate(old_core_engines): + if cur_dp_rank < new_data_parallel_size: + coro = self._call_utility_async( + "commit_prepared_elastic_ep", engine=engine + ) + else: + reconfig_request = self._make_reconfig_request( + new_data_parallel_size, + ReconfigureRankType.SHUTDOWN_CURRENT_RANK, + ) + coro = self._call_utility_async( + "reinitialize_distributed", reconfig_request, engine=engine ) - coro = self._call_utility_async( - "reinitialize_distributed", reconfig_request, engine=engine - ) reconfig_futures.append(asyncio.create_task(coro)) - # NOTE(yongji): Immediately stop sending requests to the removing engines. - self.core_engines = self.core_engines[:new_data_parallel_size] - self.lb_engines = self.lb_engines[:new_data_parallel_size] - wait_future = self._eep_wait_for_setup_switch_complete() + try: + await asyncio.gather(*reconfig_futures) - await asyncio.gather(*reconfig_futures) + self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + self._ensure_stats_update_task() + scale_down_marker = msgspec.msgpack.encode( + ("SCALE_ELASTIC_EP", new_data_parallel_size) + ) + await self.first_req_send_socket.send(scale_down_marker) + await wait_future + await self.resume_scheduler_async() + except Exception: + wait_future.cancel() + raise - self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - self._ensure_stats_update_task() - scale_down_marker = msgspec.msgpack.encode( - ("SCALE_ELASTIC_EP", new_data_parallel_size) - ) - await self.first_req_send_socket.send(scale_down_marker) - - # NOTE(yongji): Unlike scaling up, - # here we don't actually need to wait for the setup switch to complete. - # We may want to remove it in the future. - await wait_future logger.info( "[Elastic EP] Scale down completed, new data parallel size: %s", new_data_parallel_size, diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index db1896b0946..9b3bea0db9c 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -821,7 +821,10 @@ class CoreEngineActorManager: return placement_groups, local_dp_ranks def scale_up_elastic_ep( - self, cur_vllm_config: VllmConfig, new_data_parallel_size: int + self, + cur_vllm_config: VllmConfig, + new_data_parallel_size: int, + num_redundant_experts: int, ) -> None: import copy @@ -864,6 +867,9 @@ class CoreEngineActorManager: if new_data_parallel_size > 1: _apply_dp_identity_suffix(dp_vllm_config, rank) dp_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size + dp_vllm_config.parallel_config.eplb_config.num_redundant_experts = ( + num_redundant_experts + ) dp_vllm_config.parallel_config.placement_group = pg # Check if this placement group is on the head node @@ -906,39 +912,18 @@ class CoreEngineActorManager: self.created_placement_groups.append(pg) self.placement_group_is_local.append(local_client) - ray.get( - [ - actor.wait_for_init.remote() - for actor in ( - self.local_engine_actors[-new_local_engines:] - if new_local_engines > 0 - else [] - ) - + self.remote_engine_actors[ - -(len(placement_groups) - new_local_engines) : - ] - ] - ) - actors = ( self.local_engine_actors[-new_local_engines:] if new_local_engines > 0 else [] ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :] + ray.get([actor.wait_for_init.remote() for actor in actors]) for actor in actors: ref = actor.run.remote() self.run_refs.append(ref) self.actor_run_ref_dict[actor] = ref - cur_vllm_config.parallel_config.data_parallel_size = new_data_parallel_size - # Update old_vllm_config with new data_parallel_size_local if any new - # local engines were added - if new_local_engines > 0: - cur_vllm_config.parallel_config.data_parallel_size_local += ( - new_local_engines - ) - def scale_down_elastic_ep( self, cur_data_parallel_size: int, new_data_parallel_size: int ) -> None: diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 4063844d469..404acd50de9 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -116,11 +116,11 @@ class Executor(ABC): raise NotImplementedError def initialize_from_config(self, kv_cache_configs: list[KVCacheConfig]) -> None: - """ - Initialize the KV caches and begin the model execution loop of the - underlying workers. - """ + """Initialize the KV caches on the underlying workers.""" self.collective_rpc("initialize_from_config", args=(kv_cache_configs,)) + + def compile_or_warm_up_model(self) -> None: + """Compile/warm up the model and capture cudagraphs on workers.""" compilation_times: list[CompilationTimes] = self.collective_rpc( "compile_or_warm_up_model" ) From 948107acf7ef8813b8ec94fff7c5ab62d4aba9ae Mon Sep 17 00:00:00 2001 From: Xin He Date: Tue, 28 Jul 2026 20:37:55 +0800 Subject: [PATCH 22/67] [Bugfix] Enhance extra_config handling for layer name suffix matching (#48589) Signed-off-by: Xin He Co-authored-by: Kunshang Ji --- tests/quantization/test_auto_round.py | 49 +++++++++++++++++++ .../layers/quantization/inc/config_parser.py | 8 +++ 2 files changed, 57 insertions(+) diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 732080fc967..59c6a1e326b 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -243,6 +243,27 @@ def test_inc_config_parser_parallel_lm_head_defaults_to_unquantized() -> None: assert layer_config.bits == 16 +def test_inc_config_parser_suffix_match_for_lm_head() -> None: + """Short extra_config key should match fully-qualified lm_head layer name.""" + layer = object.__new__(ParallelLMHead) + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + + layer_config = config.config_parser.resolve(layer, "model.language_model.lm_head") + + assert layer_config.quantized is True + assert layer_config.bits == 4 + assert layer_config.group_size == 128 + assert layer_config.sym is True + + def test_inc_config_parser_fused_moe_requires_consistent_configs() -> None: config = make_config( extra_config={ @@ -790,6 +811,34 @@ def test_inc_get_quant_method_linear_uses_resolved_scheme(monkeypatch) -> None: assert method is sentinel +def test_inc_get_quant_method_lm_head_uses_suffix_match(monkeypatch) -> None: + """lm_head extra_config should apply to fully-qualified prefix.""" + config = make_config( + extra_config={ + "lm_head": { + "bits": 4, + "group_size": 128, + "sym": True, + } + } + ) + layer = object.__new__(ParallelLMHead) + sentinel = object() + + class DummyScheme: + def get_linear_method(self, _config, _layer, _prefix, _layer_config): + return sentinel + + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.factory.resolve_scheme", + lambda _layer_config: DummyScheme(), + ) + + method = config.get_quant_method(layer, "model.language_model.lm_head") + + assert method is sentinel + + def test_inc_get_quant_method_moe_uses_resolved_scheme(monkeypatch) -> None: config = make_config() layer = object.__new__(RoutedExperts) diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index 603b80b7cd0..6e94cad2cc6 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -142,6 +142,14 @@ class INCConfigParser: if self._config.extra_config and layer_name in self._config.extra_config: return get_config(layer_name) + # Suffix match: handle cases where extra_config keys use short names + # (e.g. "lm_head") but the layer_name is fully qualified + # (e.g. "model.language_model.lm_head") due to model nesting. + if self._config.extra_config: + for cfg_key in self._config.extra_config: + if layer_name.endswith(f".{cfg_key}"): + return get_config(cfg_key) + quantized = not isinstance(layer, ParallelLMHead) if self._config.block_name_to_quantize: quantized = any( From 98e91a9600eb75b2de14ef27f13b10088d1a1279 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 28 Jul 2026 14:57:12 +0200 Subject: [PATCH 23/67] [PD][NixlPush] Skip extra `add_remote_agent` step in D->P handshake (#49345) Signed-off-by: NickLucche --- .../kv_transfer/kv_connector/v1/nixl/base_worker.py | 3 ++- .../kv_transfer/kv_connector/v1/nixl/push_worker.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index d4e516d6df7..9063308ce2e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -700,8 +700,9 @@ class NixlBaseConnectorWorker: ) setup_agent_time = time.perf_counter() logger.debug( - "NIXL handshake: add agent took: %s", + "NIXL handshake: add agent took: %s (notif_agents_only=%s)", setup_agent_time - got_metadata_time, + notif_agents_only, ) remote_ranks = (remote_pp_rank, remote_rank) remote_rank_to_agent_name[remote_ranks] = remote_agent_name diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 43b4f893907..8cb710e7e58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -299,8 +299,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): reg_data["remote_port"], reg_data["remote_tp_size"], pp_size=remote_pp_size, - # D never addresses P memory in push mode; just load P's agents. - notif_agents_only=remote_pp_size > 1, + # D only ever sends PUSH_REG notifs to P and never reads or writes + # P's memory in push mode, so it never needs the transfer + # descriptors set up by the full add_remote_agent path. + notif_agents_only=True, ) if fut is None: self._do_send_reg_notif(req_id, reg_data) From 9b9fc4039c25a6e4fe0ae97361b62edd74b8b47e Mon Sep 17 00:00:00 2001 From: liangel-02 Date: Tue, 28 Jul 2026 07:26:55 -0600 Subject: [PATCH 24/67] add epilogue hook to flex attention (#45841) Signed-off-by: Angel Li Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Michael Goin --- vllm/v1/attention/backends/flex_attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index e7fe5852194..83144751aeb 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -12,6 +12,7 @@ import torch import torch._dynamo.decorators import torch.nn.functional as F from torch.nn.attention.flex_attention import ( + AuxRequest, BlockMask, _mask_mod_signature, _score_mod_signature, @@ -1228,6 +1229,9 @@ class FlexAttentionImpl(AttentionImpl): if block_n is not None: self.block_n = block_n + # Optional post-attention epilogue transform + self.out_transform = kwargs.get("out_transform") + @staticmethod def view_as_4d(tensor: torch.Tensor) -> torch.Tensor: """View a 3d tensor as 4D.""" @@ -1392,8 +1396,13 @@ class FlexAttentionImpl(AttentionImpl): self.scale, enable_gqa=enable_gqa, kernel_options=kernel_options, + return_aux=AuxRequest(lse=True) if self.out_transform is not None else None, ) + if self.out_transform is not None: + out, aux = out + out = self.out_transform(out, aux.lse) + # Flex doesn't have an out variant today, rely on epilogue fusion out = out.permute(0, 2, 1, 3).squeeze(0) output[:num_actual_tokens, :, :].copy_(out) From 601fa9a74e8cd2ae0a4fa127d68b91bbf363a24e Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 06:45:02 -0700 Subject: [PATCH 25/67] [KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models (#49612) Signed-off-by: Nick Hill Co-authored-by: Claude Fable 5 --- .../kv_connector/unit/test_nixl_connector.py | 10 +- .../unit/test_nixl_connector_hma.py | 114 ++++ .../unit/test_nixl_desc_geometry.py | 619 ++++++++++++++++++ tests/v1/kv_connector/unit/test_tp_mapping.py | 54 ++ .../kv_connector/v1/nixl/base_worker.py | 271 ++++++-- .../kv_connector/v1/nixl/pull_worker.py | 35 +- .../kv_connector/v1/nixl/push_worker.py | 22 +- 7 files changed, 1014 insertions(+), 111 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_nixl_desc_geometry.py diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d58f117fa31..c474077ab22 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout - # Mock register_kv_caches attribute needed for tests that do not call it. + # Mock register_kv_caches attributes needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} + self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) @@ -765,8 +766,9 @@ class TestNixlHandshake: assert remote_info.remote_tp_size == remote_tp_size assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size) # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks - assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio - assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio + split_key = (-tp_ratio, worker.block_size) + assert split_key in worker.src_xfer_handles_by_tp_ratio + assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio assert remote_engine_id in worker.dst_xfer_side_handles assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set( range(tp_ratio) @@ -2091,7 +2093,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 4945942ba3a..1f7a62d2c9a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch(): assert list(result) == expected, f"Expected {expected}, got {list(result)}" +@pytest.mark.cpu_test +def test_get_block_descs_ids_hetero_block_size_hybrid(): + """With a block-size ratio, FA desc ids are ratio-expanded while SSM + desc ids keep the unexpanded logical stride (state blocks are never + sub-split).""" + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = _make_mock_worker_for_desc_ids( + num_regions=2, + has_mamba=True, + group_spec_types=(FullAttentionSpec, MambaSpec), + block_len_per_layer=[100], + ) + + ratio = 4 + # FA ids are already remote-granularity (expanded) sub-block ids. + fa_sub_blocks = [3, 5] + ssm_blocks = [1] + result = worker._compute_desc_ids( + block_ids=(fa_sub_blocks, ssm_blocks), + dst_num_blocks=100, + block_size_ratio=ratio, + physical_blocks_per_logical=1, + ) + + # FA regions have 100*4 entries each; SSM regions (4 per layer) start at + # 2*400 and stride by the unexpanded 100 logical blocks. + expected = [3, 5, 403, 405, 801, 901, 1001, 1101] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +def _bind_worker_method(worker, name): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + method = getattr(NixlConnectorWorker, name) + setattr(worker, name, method.__get__(worker, NixlConnectorWorker)) + + +@pytest.mark.cpu_test +def test_map_block_ids_for_block_size_ratio_hybrid(): + """Attention groups expand to remote granularity and clip to the remote + coverage; mamba state blocks pass through 1:1.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + _bind_worker_method(worker, "get_mapped_blocks") + _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio") + + local, remote = worker._map_block_ids_for_block_size_ratio( + [[1, 2, 3], [7]], + [list(range(30, 40)), [42]], + 4, + ) + # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks. + assert local == [list(range(4, 14)), [7]] + assert remote == [list(range(30, 40)), [42]] + + # Attention-only full prefix hit: empty local list is preserved. + worker._group_spec_types = (FullAttentionSpec,) + local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4) + assert local == [] + + +@pytest.mark.cpu_test +def test_post_process_zeroes_untransferred_tail(): + """The untransferred sub-blocks of the last local block are zeroed on + receive; mamba state caches are untouched by the attention permute.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + ratio = 4 + block_tokens = 8 # 2 tokens per remote sub-block + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.device_type = "cpu" + worker.enable_permute_local_kv = False + attn_cache = torch.ones(6, block_tokens, 2, 4) + mamba_cache = torch.ones(6, 16) + worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + fa_group = MagicMock(layer_names=["attn.0"]) + ssm_group = MagicMock(layer_names=["mamba.0"]) + worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group]) + # The cached property filters mamba layers out of the permuted caches. + attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) + assert len(attn_caches) == 1 and attn_caches[0] is attn_cache + worker._attention_kv_caches = attn_caches + _bind_worker_method(worker, "post_process_device_kv_on_receive") + + # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + + # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). + assert torch.all(attn_cache[2] == 1) + assert torch.all(attn_cache[3, :4] == 1) + assert torch.all(attn_cache[3, 4:] == 0) + # Untouched blocks and the mamba cache keep their content. + assert torch.all(attn_cache[4] == 1) + assert torch.all(mamba_cache == 1) + + @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py new file mode 100644 index 00000000000..e2fc41a0229 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models +under heterogeneous P/D block geometry (TP-sharded KDA-style state, so +the mamba-aligned logical block size differs between P and D while the +kernel-granularity pages stay equal). + +The invariant under test: every LOCAL byte range a request's READ transfers +into must lie within that request's own blocks. A violation means an +incoming transfer can overwrite a co-resident request's KV or mamba state +mid-decode (silent corruption of an unrelated request). +""" + +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from .utils import create_vllm_config + + +class _RecordingNixl: + """Minimal NIXL wrapper stand-in that records descriptor lists and + prepared transfers so tests can resolve desc ids to byte ranges.""" + + def __init__(self, *args, **kwargs): + self.dlists: dict[int, np.ndarray] = {} + self.xfers: list[tuple] = [] + self._next_handle = 1 + + def get_reg_descs(self, caches_data, mem_type): + return caches_data + + def register_memory(self, descs, backends=None): + pass + + def deregister_memory(self, descs): + pass + + def get_agent_metadata(self): + return b"agent-meta" + + def get_xfer_descs(self, blocks_data, mem_type): + return blocks_data + + def prep_xfer_dlist(self, agent, descs): + handle = self._next_handle + self._next_handle += 1 + self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3) + return handle + + def add_remote_agent(self, metadata): + return "remote-agent" + + def make_prepped_xfer( + self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None + ): + handle = self._next_handle + self._next_handle += 1 + self.xfers.append( + ( + op, + local_handle, + np.asarray(local_ids), + remote_handle, + np.asarray(remote_ids), + ) + ) + return handle + + def transfer(self, handle): + pass + + def check_xfer_state(self, handle): + return "DONE" + + def get_xfer_telemetry(self, handle): + from types import SimpleNamespace + + return SimpleNamespace( + xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1 + ) + + def release_xfer_handle(self, handle): + pass + + def release_dlist_handle(self, handle): + pass + + def send_notif(self, agent, notif_msg=None): + pass + + def get_new_notifs(self): + return {} + + def remove_remote_agent(self, agent): + pass + + +def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks): + """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( + base_worker as bw, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + mla_spec = MLAAttentionSpec( + block_size=local_block_size, + num_kv_heads=1, + head_size=6, + dtype=torch.float16, + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=local_block_size, + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_logical_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_logical_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + vllm_config = create_vllm_config(block_size=local_block_size) + vllm_config.cache_config.enable_prefix_caching = False + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared-tensor regions this test builds + # would not be deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + from unittest.mock import MagicMock + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + from vllm.config import set_current_vllm_config + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + worker.use_mla = True + + # Attention caches are kernel-block granular on dim 0, as the + # receive post-process assumes. + ppl = local_block_size // kernel_block_size + tensors = [ + torch.zeros( + num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8 + ) + for _ in range(2) + ] + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + # Keep tensors alive alongside the worker; flat views for byte checks. + worker._test_tensors = [t.view(-1) for t in tensors] + worker._test_tensors_2d = tensors + worker._test_unified_page = unified_page + return worker + + +def _make_remote_meta( + worker, + remote_block_size, + remote_kernel_block_size, + remote_num_logical, + remote_ssm_sizes, +): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, + ) + + remote_ppl = remote_block_size // remote_kernel_block_size + # Kernel-granularity pages are TP-independent for MLA hybrids and must + # match the local ones for the handshake to pass, scaled down by the + # block-size ratio when the remote's kernel block is smaller. + block_size_ratio = worker.block_size // remote_kernel_block_size + kernel_page = worker.block_len_per_layer[0] // block_size_ratio + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"remote-agent-meta", + device_id=0, + kv_caches_base_addr=[0x10_000_000, 0x20_000_000], + num_blocks=remote_num_logical * remote_ppl, + block_lens=[kernel_page, kernel_page], + kv_cache_layout=worker.kv_cache_layout, + block_size=remote_kernel_block_size, + ssm_sizes=remote_ssm_sizes, + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=remote_ppl, + ) + + +def _owned_byte_ranges(worker, group_logical_ids): + """Byte ranges owned by a request: for each HMA region tensor, every + logical block id of every group maps to one unified page.""" + unified_page = worker._test_unified_page + bases = [t.data_ptr() for t in worker._test_tensors] + owned = [] + for base in bases: + for ids in group_logical_ids: + for b in ids: + owned.append((base + b * unified_page, base + (b + 1) * unified_page)) + return owned + + +def _assert_local_writes_within(worker, owned_ranges): + nixl = worker.nixl_wrapper + assert nixl.xfers, "no transfers were posted" + violations = [] + total_descs = 0 + for op, local_handle, local_ids, _, remote_ids in nixl.xfers: + assert len(local_ids) == len(remote_ids) + desc_arr = nixl.dlists[local_handle] + for i in local_ids: + addr, length, _dev = desc_arr[int(i)] + addr, length = int(addr), int(length) + total_descs += 1 + if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges): + violations.append((int(i), hex(addr), length)) + assert not violations, ( + f"{len(violations)}/{total_descs} local descriptors write outside " + f"the request's own blocks: {violations[:10]}" + ) + return total_descs + + +@pytest.mark.cpu_test +def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): + """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens + (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), + equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and + TP-sharded KDA state. Every local descriptor of the request's reads must + stay within its own blocks.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + assert worker._physical_blocks_per_logical_kv_block == 3 + + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + for rank in (0, 1): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + + # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok + # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. + # Sparse, non-contiguous ids so neighbor blocks exist on all sides. + local_ids = ([2, 5], [1], [7]) + remote_ids = [[1, 4], [5], [2]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 2, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + worker._read_blocks_for_req("req-b", meta) + + owned = _owned_byte_ranges(worker, local_ids) + total = _assert_local_writes_within(worker, owned) + # Multi-read: rank 0 carries the replicated MLA + its SSM shard, + # rank 1 carries only its SSM shard. + assert len(worker.nixl_wrapper.xfers) == 2 + assert total > 0 + + +def _resolve( + desc_arr, + idx, + bases, + region_size, + unified_page, + desc_page, + logical_ids_attn, + block_tokens, +): + """Resolve a desc id to (region, kind, token_start) where kind is 'attn' + (desc-page sized, sub-block-aligned, in the request's attention blocks) + or 'mamba'. token_start is the request-relative token offset, so local + and remote are comparable even when their kernel blocks differ in size.""" + addr, length, _ = (int(x) for x in desc_arr[int(idx)]) + for region, base in enumerate(bases): + off = addr - base + if 0 <= off < region_size: + b = off // unified_page + rem = off % unified_page + if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn: + pos = logical_ids_attn.index(b) + tokens_per_desc = block_tokens * desc_page // unified_page + sub = rem // desc_page + return (region, "attn", pos * block_tokens + sub * tokens_per_desc) + return (region, "mamba", None) + raise AssertionError(f"desc {idx} addr {addr:#x} not in any region") + + +def _run_hetero_case( + local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None +): + """Full pull-path run for one geometry; returns pairing records. + + ``remote_kernel`` defaults to the local kernel block size; a smaller + value additionally exercises block_size_ratio > 1. + """ + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + remote_kernel = remote_kernel or kernel + block_size_ratio = kernel // remote_kernel + remote_ppl = remote_block // remote_kernel + matched = num_tokens - 1 # mamba N-1 rule + n_local = -(-num_tokens // local_block) + n_remote = -(-matched // remote_block) + + worker = _make_mla_hybrid_worker( + local_block_size=local_block, + kernel_block_size=kernel, + num_logical_blocks=max(2 * n_local + 4, 8), + ) + # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size + # shards of each. + meta_r = _make_remote_meta( + worker, + remote_block_size=remote_block, + remote_kernel_block_size=remote_kernel, + remote_num_logical=max(2 * n_remote + 4, 8), + remote_ssm_sizes=(48 // tp_size, 64 // tp_size), + ) + for rank in range(tp_size): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + + # Sparse ids so neighbors exist between the request's blocks. + local_attn = [2 * i + 1 for i in range(n_local)] + remote_attn = [2 * i + 2 for i in range(n_remote)] + local_ids = (local_attn, [0], [2 * n_local + 2]) + remote_ids = [remote_attn, [1], [0]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": tp_size, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + # Sentinel-fill the local KV so untouched bytes are detectable. + for t in worker._test_tensors: + t.fill_(0xAA) + + worker._read_blocks_for_req("req-b", meta) + + # Invariant 1: all local writes within the request's own blocks. + owned = _owned_byte_ranges(worker, local_ids) + _assert_local_writes_within(worker, owned) + + # Invariant 2: local<->remote attention pairs are token-aligned. + nixl = worker.nixl_wrapper + local_bases = [t.data_ptr() for t in worker._test_tensors] + remote_bases = [0x10_000_000, 0x20_000_000] + local_unified = worker._test_unified_page + remote_unified = (local_unified // local_block) * remote_block + # With block_size_ratio > 1 the local page is split into ratio sub-descs, + # each the size of a whole remote kernel page. + desc_page = worker.block_len_per_layer[0] // block_size_ratio + meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified + covered_tokens = set() + for op, lh, lids, rh, rids in nixl.xfers: + larr, rarr = nixl.dlists[lh], nixl.dlists[rh] + for li, ri in zip(lids, rids): + lreg, lkind, ltok = _resolve( + larr, + li, + local_bases, + len(worker._test_tensors[0]), + local_unified, + desc_page, + local_attn, + local_block, + ) + rreg, rkind, rtok = _resolve( + rarr, + ri, + remote_bases, + meta_r_num_blocks_bytes, + remote_unified, + desc_page, + remote_attn, + remote_block, + ) + assert lkind == rkind, ( + f"pair kind mismatch: local {lkind} vs remote {rkind} " + f"(local desc {li}, remote desc {ri})" + ) + assert lreg == rreg, ( + f"region mismatch: local {lreg} vs remote {rreg} for " + f"tokens {ltok} vs {rtok}" + ) + if lkind == "attn": + assert ltok == rtok, ( + f"TOKEN MISALIGNMENT: local sub-block holds tokens " + f"[{ltok}..) but receives remote tokens [{rtok}..) " + f"(geometry local_block={local_block}, " + f"remote_block={remote_block}, N={num_tokens})" + ) + covered_tokens.add(ltok) + + # Invariant 3: full coverage of the matched tokens, at the finest + # transfer granularity (the remote kernel block). + needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)} + missing = needed - covered_tokens + assert not missing, ( + f"tokens never transferred: {sorted(missing)[:8]} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + # Invariant 4: no stale bytes after receive completion. The scheduler + # excludes the blocks covering the matched tokens from alloc-time KV + # zeroing (the zeroing would race the RDMA write), so every byte of + # those blocks must be either written by the transfer or zeroed by the + # receive post-process. Stale bytes surface as mid-response garbage + # once decode grows into the untransferred tail. + for op, lh, lids, rh, rids in nixl.xfers: + larr = nixl.dlists[lh] + for li in lids: + addr, length, _ = (int(x) for x in larr[int(li)]) + for t in worker._test_tensors: + off = addr - t.data_ptr() + if 0 <= off < t.numel(): + t[off : off + length] = 0 # simulate the RDMA write + break + done_sending, done_recving = worker.get_finished() + assert "req-b" in done_recving + n_excluded = -(-matched // local_block) + stale = [] + for b in local_attn[:n_excluded]: + for region, t in enumerate(worker._test_tensors): + page = t[b * local_unified : (b + 1) * local_unified] + n_stale = int((page == 0xAA).sum()) + if n_stale: + stale.append((region, b, n_stale)) + assert not stale, ( + f"stale (unzeroed, untransferred) bytes in matched-range attention " + f"blocks (region, block, bytes): {stale} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "local_block,remote_block", + [ + (12, 8), # ppl 3 vs 2 + (36, 8), # ppl 9 vs 2 (large ppl asymmetry, scaled) + (24, 4), # ppl 6 vs 1 + (16, 24), # remote larger than local (D_TP > P_TP direction) + ], +) +@pytest.mark.parametrize("num_tokens", list(range(2, 40))) +def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): + """Sweep prompt lengths across block-boundary residues for several + hetero-ppl geometries; assert neighbor-safety, token alignment, and + coverage of every transferred kernel block.""" + _run_hetero_case( + local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around the remote kernel block (4), the local kernel block + # (8), the remote logical block (8) and the local logical block (24). + [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49], +) +def test_hetero_ppl_with_block_size_ratio(num_tokens): + """Both hetero regimes at once: kernel blocks differ (local 8 / remote + 4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs + 2). The transfer is clipped at remote sub-block granularity by the + pairing and front-trimmed by _apply_prefix_caching, so the + untransferred tail can span both a partial block and whole blocks — + the case each of the two former zeroing paths handled only half of.""" + _run_hetero_case( + local_block=24, + kernel=8, + remote_block=8, + remote_kernel=4, + num_tokens=num_tokens, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around every geometric boundary: kernel block (64), remote + # logical block (768), local logical block (5760), plus odd offsets. + [ + 2, + 63, + 64, + 65, + 127, + 128, + 300, + 640, + 767, + 768, + 769, + 831, + 832, + 1000, + 1535, + 1536, + 1537, + 2303, + 2304, + 2305, + 3001, + 5759, + 5760, + 5761, + 5824, + 6528, + 6529, + ], +) +def test_mla_hybrid_large_ppl_geometry(num_tokens): + """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): + decode (local) logical block 5760 / kernel 64 (ppl=90), prefill + (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with + replicated MLA and 8-way TP-sharded KDA state.""" + _run_hetero_case( + local_block=5760, + kernel=64, + remote_block=768, + num_tokens=num_tokens, + tp_size=8, + ) diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 7c735098230..d72bbc9ec2d 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles: # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0) # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0) assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)] + + def test_hetero_block_size_splits(self): + """With a block-size ratio, single-source FA sub-block descs pass + through whole; SSM descs are unexpanded and split per source.""" + plan = TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded). + src_blocks_data = np.array( + [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + (3000, 400, 0), + ], + dtype=np.uint64, + ) + + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2)) + + assert len(splits) == 2 + fa_passthrough = [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + ] + assert splits[0] == fa_passthrough + [(3000, 200, 0)] + assert splits[1] == fa_passthrough + [(3200, 200, 0)] + + def test_hetero_block_size_head_sharded_asserts(self): + """Head-sharded FA reads (multiple FA sources) are incompatible with + a block-size mismatch and must fail loudly.""" + plan = TPMapping( + source_ranks_per_group=((0, 1), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 1}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + src_blocks_data = np.array( + [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)], + dtype=np.uint64, + ) + + with pytest.raises(AssertionError, match="Head-sharded"): + list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 9063308ce2e..480ac8937df 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -12,6 +12,7 @@ import uuid from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor +from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -67,6 +68,7 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -122,9 +124,11 @@ class NixlBaseConnectorWorker: return (region_ids * num_blocks + block_arr).flatten() # Compute desc ids per group using the right stride: FA descs have - # num_blocks entries per region (kernel granularity), SSM descs have - # logical_blocks entries per region (no kernel splitting). - logical_blocks = num_blocks // physical_blocks_per_logical + # num_blocks entries per region (kernel granularity, expanded by + # block_size_ratio for heterogeneous block sizes), SSM descs have + # logical_blocks entries per region (no kernel splitting, and never + # ratio-expanded since state blocks are indivisible). + logical_blocks = dst_num_blocks // physical_blocks_per_logical all_descs: list[np.ndarray] = [] for i, group in enumerate(block_ids): group_arr = np.asarray(group) @@ -161,6 +165,7 @@ class NixlBaseConnectorWorker: plan: TPMapping, src_blocks_data: np.ndarray, num_fa_descs: int, + block_size_ratio: int = 1, ) -> Iterator[list[tuple[int, int, int]]]: """Build split handle data for P_TP > D_TP scenario. @@ -187,6 +192,11 @@ class NixlBaseConnectorWorker: # Per-FA-descriptor replicate flag, in _build_fa_local emission order. fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + + assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), ( + "Head-sharded attention reads with P_TP > D_TP and heterogeneous " + "block sizes are not supported" + ) src_blocks_list = src_blocks_data.tolist() for p_idx, p_rank in enumerate(plan.all_source_ranks): @@ -442,9 +452,12 @@ class NixlBaseConnectorWorker: # nixl_prepped_dlist_handle. self.src_xfer_handles_by_block_size: dict[int, int] = {} + # Local descriptor arrays per remote block size (block_size_ratio>1), + # kept for building per-tp-ratio splits at the same granularity. + self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {} # Populated dynamically during handshake based on remote configuration. - # Keep track of regions at different tp_ratio values. tp_ratio->handles - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} + # Per-source split handles, keyed by (tp_ratio, remote_block_size). + self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {} # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) @@ -1258,11 +1271,7 @@ class NixlBaseConnectorWorker: agent_metadata_bytes=encoder.encode(agent_metadata), ) - def _build_mamba_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> np.ndarray: + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for local mamba blocks with DS conv layout, as an Nx3 uint64 array. @@ -1289,16 +1298,17 @@ class NixlBaseConnectorWorker: | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | +-------------------+ +--------------------+ |1st_split-2nd_split| |1st_split-2nd_split | + + Mamba state blocks are indivisible (not token-extent data), so the + descriptors always use the local page geometry regardless of any + attention block-size ratio; their desc ids are likewise never + ratio-expanded (see _compute_desc_ids). """ - assert block_size_ratio == 1, ( - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " - f"Got block_size_ratio={block_size_ratio}." - ) assert base_addresses, "Local KV cache base addresses must not be empty." assert self._conv_decomp is not None conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size - num_blocks = self._logical_num_blocks * block_size_ratio + num_blocks = self._logical_num_blocks physical_per_logical = self._physical_blocks_per_logical_kv_block device_id = self.device_id block_arange = np.arange(num_blocks, dtype=np.uint64) @@ -1307,9 +1317,7 @@ class NixlBaseConnectorWorker: for i, base_addr in enumerate(base_addresses): # Jump one page_size, but ssm page_size may be bigger when kernel # locks block size to a specific value (physical_per_logical scale). - page_stride = ( - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical - ) + page_stride = self.block_len_per_layer[i] * physical_per_logical blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) @@ -1459,7 +1467,7 @@ class NixlBaseConnectorWorker: self.device_id, ) if self._has_mamba: - assert self.num_descs == len(blocks_data) + assert self.num_descs * block_size_ratio == len(blocks_data) # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split # is unnecessary — a single conv desc per block suffices. Consider # adding a fast path that falls back to the standard 2-region @@ -1467,7 +1475,7 @@ class NixlBaseConnectorWorker: # remote has been seen. Currently we always register 4 regions # because local descs are created before knowing the remote TP. logger.debug("Registering local Mamba descriptors (4 regions/layer)") - mamba = self._build_mamba_local(local_base_addresses, block_size_ratio) + mamba = self._build_mamba_local(local_base_addresses) blocks_data = np.concatenate([blocks_data, mamba]) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1607,29 +1615,44 @@ class NixlBaseConnectorWorker: plan = self.tp_mappings[engine_id] + ### (Optional) Register a local handler at the remote engine's block + ### granularity (remote/prefill blocks smaller than local). + remote_block_size = nixl_agent_meta.block_size + src_blocks_data = self.src_blocks_data + if block_size_ratio > 1: + if remote_block_size not in self.src_xfer_handles_by_block_size: + handle, blocks_data = self.register_local_xfer_handler( + remote_block_size + ) + self.src_xfer_handles_by_block_size[remote_block_size] = handle + self.src_blocks_data_by_block_size[remote_block_size] = blocks_data + src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size] + ### (Optional) Register local agent memory regions. MLA is not split. + split_key = (tp_ratio, remote_block_size) if ( tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1) - and tp_ratio not in self.src_xfer_handles_by_tp_ratio + and split_key not in self.src_xfer_handles_by_tp_ratio ): # Remote tp_size > local tp_size: read from multiple remote ranks. # Logically "split" own regions into per-source chunks. Hybrid # MLA+SSM also needs this path: MLA is replicated and read once, # while the SSM state is sharded across every remote TP rank. - # We only do this once per remote tp_size (replica-friendly). - self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] + # We only do this once per remote (tp_size, block_size). + self.src_xfer_handles_by_tp_ratio[split_key] = [] for handle_data in self._build_local_splits_from_plan( plan, - self.src_blocks_data, - self.num_descs, + src_blocks_data, + self.num_descs * block_size_ratio, + block_size_ratio, ): descs = self.nixl_wrapper.get_xfer_descs( handle_data, self.nixl_memory_type ) handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + self.src_xfer_handles_by_tp_ratio[split_key].append(handle) ### Register remote agent memory regions # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With @@ -1665,13 +1688,6 @@ class NixlBaseConnectorWorker: self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) ) - if block_size_ratio > 1: - # when prefill with smaller block_size, we need to init a - # new handler with same block_len to match - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] - ) - return remote_agent_name def _validate_remote_agent_handshake( @@ -1716,9 +1732,13 @@ class NixlBaseConnectorWorker: "Disable prefix caching with --no-enable-prefix-caching." ) - if self._is_hma_required: - assert block_size_ratio == 1, ( - "HMA does not support different remote block size yet" + if block_size_ratio != 1: + # Heterogeneous block sizes transfer at remote-block granularity; + # the untransferred tail of the last local attention block is + # zeroed in the receive post-process, and mamba state pages + # transfer 1:1 (never sub-split). + assert not self.use_host_buffer, ( + "Heterogeneous block sizes are not supported with host buffer" ) kv_cache_layout = ( self.kv_cache_layout @@ -1875,27 +1895,57 @@ class NixlBaseConnectorWorker: "d2h", ) + @cached_property + def _attention_kv_caches(self) -> list[torch.Tensor]: + """Device KV caches of attention layers (mamba states excluded), + as consumed by the receive post-process.""" + assert self.device_kv_caches, ( + "_attention_kv_caches accessed before register_kv_caches" + ) + mamba_layers = { + name + for g, group in enumerate(self.kv_cache_config.kv_cache_groups) + if _is_ssm_spec(self._group_spec_types[g]) + for name in group.layer_names + } + kv_caches = self.device_kv_caches + return [cache for name, cache in kv_caches.items() if name not in mamba_layers] + def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[list[int]], + block_ids_list: list[tuple[list[int], int]], + convert: bool = True, ): """ Post process device kv cache after receiving from remote. - 3 types of post processing supported: + 3 types of conversion supported (``convert``): * kv_cache_postprocess_layout => convert from HND to NHD * kv_cache_postprocess_blksize => convert from small block size to large block size * kv_cache_postprocess_blksize_and_layout => convert from small block size to large block size and convert from HND to NHD + The transfer only covers ``covered_sub_blocks`` remote-sized + sub-blocks of each request's local attention blocks; the rest was + clipped, either by remote-block pairing (block-size ratio) or by the + hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were + excluded from the scheduler's alloc-time KV zeroing (which would race + the RDMA write), so everything past the covered range is zeroed here. + Stale bytes would otherwise surface as garbage or NaNs once decode + grows into the untransferred tail. """ if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." assert self.transfer_topo is not None - if self.enable_permute_local_kv and block_size_ratio > 1: + if not convert: + logger.debug( + "Post-processing device kv cache on receive by zeroing " + "untransferred blocks." + ) + elif self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " "block_size with %sx bigger and permuting layout from HND" @@ -1914,18 +1964,45 @@ class NixlBaseConnectorWorker: block_size_ratio, ) - for block_ids in block_ids_list: - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + attn_caches = self._attention_kv_caches + device = attn_caches[0].device + for block_ids, covered_sub_blocks in block_ids_list: + # Blocks the transfer didn't write: the token tail of the last + # partially covered block, then everything beyond it. + covered_blocks, sub_blocks_in_last = divmod( + covered_sub_blocks, block_size_ratio + ) + first_stale = covered_blocks + (1 if sub_blocks_in_last else 0) + has_stale = first_stale < len(block_ids) + indices = None + if convert or has_stale: + indices = async_tensor_h2d(block_ids, device, torch.long) - for cache in self.device_kv_caches.values(): - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive( - cache, indices, block_size_ratio - ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) + if convert: + for cache in attn_caches: + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive( + cache, indices, block_size_ratio + ) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive( + cache, indices, block_size_ratio + ) + + if sub_blocks_in_last: + last_block_id = block_ids[covered_blocks] + for cache in attn_caches: + # Both post-processed layouts leave tokens on dim 1. + sub_block_tokens = cache.shape[1] // block_size_ratio + zero_from = sub_blocks_in_last * sub_block_tokens + cache[last_block_id, zero_from:].zero_() + if has_stale: + assert indices is not None + stale_ids = indices[first_stale:] + for cache in attn_caches: + cache.index_fill_(0, stale_ids, 0) def post_process_device_kv_on_receive_heterogeneous_attn( self, block_ids: list[int] @@ -1995,18 +2072,33 @@ class NixlBaseConnectorWorker: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) - # post processing for heteroblocksize + # Post processing for heteroblocksize/layout, and for blocks the + # transfer clipped. The latter happens either at remote-block + # granularity (block_size_ratio > 1) or at kernel-block + # granularity, when equal kernel pages meet differing logical + # block sizes and _apply_prefix_caching front-trims to the + # minimum count (hybrid heterogeneous TP). remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) block_size_ratio = self.transfer_topo.block_size_ratio( remote_info.remote_block_size ) - if not self.use_mla and ( - block_size_ratio > 1 or self.enable_permute_local_kv - ): - assert not self._is_hma_required - block_ids_for_blocksize_post_process[block_size_ratio].append( - meta.local_physical_block_ids[0] - ) + hetero_ppl = ( + remote_info.remote_physical_blocks_per_logical + != self._physical_blocks_per_logical_kv_block + ) + if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl: + for g, local_group in enumerate(meta.local_physical_block_ids): + if not local_group or _is_ssm_spec(self._group_spec_types[g]): + continue + # Number of remote-sized sub-blocks the transfer covered; + # everything past this was clipped and must be zeroed. + covered_sub_blocks = min( + len(local_group) * block_size_ratio, + len(meta.remote.block_ids[g]), + ) + block_ids_for_blocksize_post_process[block_size_ratio].append( + (local_group, covered_sub_blocks) + ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: block_ids_for_heterogeneous_attn_post_process.append( @@ -2016,7 +2108,14 @@ class NixlBaseConnectorWorker: block_size_ratio, block_ids_list, ) in block_ids_for_blocksize_post_process.items(): - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + # MLA never needs the block-size/layout conversion, but its + # clipped blocks still need zeroing. + convert = not self.use_mla and ( + block_size_ratio > 1 or self.enable_permute_local_kv + ) + self.post_process_device_kv_on_receive( + block_size_ratio, block_ids_list, convert + ) for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) @@ -2206,6 +2305,45 @@ class NixlBaseConnectorWorker: return mapped_2d.flatten().astype(np.int64) + def _map_block_ids_for_block_size_ratio( + self, + local_block_ids: BlockIds, + remote_block_ids: BlockIds, + block_size_ratio: int, + ) -> tuple[BlockIds, BlockIds]: + """Map attention-group block ids to remote-block granularity. + + Each local attention block is split into ``block_size_ratio`` + sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the + remote list — the untransferred tail of the last local block — are + clipped here and zeroed in the receive post-process. Mamba state + blocks are indivisible and transfer 1:1, unexpanded. + + ex: remote (prefill) block ids with block_size 4: + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + Local (decode) block ids with block_size 16: [1, 2, 3] expand to + [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote. + """ + mapped_local: list[list[int]] = [] + mapped_remote: list[list[int]] = [] + for i, remote_group in enumerate(remote_block_ids): + local_group = local_block_ids[i] if local_block_ids else [] + if _is_ssm_spec(self._group_spec_types[i]): + mapped_local.append(list(local_group)) + mapped_remote.append(list(remote_group)) + continue + mapped = self.get_mapped_blocks( + np.asarray(local_group), block_size_ratio + ).tolist() + if len(mapped) > len(remote_group): + mapped = mapped[: len(remote_group)] + mapped_local.append(mapped) + mapped_remote.append(list(remote_group)) + if not any(mapped_local): + # Full prefix cache hit is indicated with an empty list. + return [], mapped_remote + return mapped_local, mapped_remote + def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: """ Convert block ids to kernel physical block ids. @@ -2300,13 +2438,18 @@ class NixlBaseConnectorWorker: remote_block_ids[i] = remote_group[-num_local_blocks:] else: # TODO Handle prefix caching with different block_sizes - max_padding = max( - self._physical_blocks_per_logical_kv_block, - remote_physical_per_logical, + # Allocation rounding legitimately leaves up to + # ppl - 1 trailing dead kernel blocks per side (plus one + # extra local block for the recomputed final token), so + # the counts may differ by up to the sum of the two + # ratios; anything larger indicates mismatched lists. + max_padding = ( + self._physical_blocks_per_logical_kv_block + + remote_physical_per_logical ) - assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( + assert abs(num_local_blocks - num_remote_blocks) <= max_padding, ( f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| >= {max_padding}" + f"{num_remote_blocks}| > {max_padding}" ) num_blocks = min(num_local_blocks, num_remote_blocks) local_block_ids[i] = local_block_ids[i][:num_blocks] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 40d6851769c..0c66596c4b7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -5,8 +5,6 @@ import time from typing import TYPE_CHECKING -import numpy as np - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -178,10 +176,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): ) # Get side handles. if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1): - assert remote_block_size == self.block_size # Remote tp_size > local tp_size: we must perform multiple # reads. Get the memory chunk onto which we will write to. - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: # Single read from remote, we write to the whole memory region. # Also handle remote block size different from local block size. @@ -235,30 +233,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - # NOTE: - # get_mapped_blocks will always expand block_ids for n times. - # ex: - # prefill block_ids with block_size as 4: - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - # Local decode block_ids with block_size as 16: [1, 2, 3] - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # Then we clip local to align with prefill - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) # NOTE(rob): having the staging blocks be on the READER side is # not going to work well (since we will have to call rearrange tensors). # after we detect the txn is complete (which means we cannot make the diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 8cb710e7e58..6bc48963c58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -39,7 +39,6 @@ from concurrent.futures import Future from typing import TYPE_CHECKING, Any import msgspec -import numpy as np from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( @@ -553,8 +552,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): req_id, ) if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: local_xfer_side_handle = self.src_xfer_handles_by_block_size[ remote_block_size @@ -606,18 +605,11 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) notif_id = f"{remote_request_id}:{self.world_size}".encode() From 62d8db7c05af8b9ef3655cf13d68416dbe3185d8 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:47:04 +0100 Subject: [PATCH 26/67] [Bugfix] Add missing `vllm/models/kimi_k3/__init__.py` (#50131) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/models/kimi_k3/__init__.py | 2 ++ vllm/models/kimi_k3/amd/ops/__init__.py | 2 ++ 2 files changed, 4 insertions(+) create mode 100644 vllm/models/kimi_k3/__init__.py diff --git a/vllm/models/kimi_k3/__init__.py b/vllm/models/kimi_k3/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/__init__.py b/vllm/models/kimi_k3/amd/ops/__init__.py index e69de29bb2d..208f01a7cb5 100644 --- a/vllm/models/kimi_k3/amd/ops/__init__.py +++ b/vllm/models/kimi_k3/amd/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project From 94100b5915d449aed0685cc5c6fc8949fcf5fe40 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 07:02:31 -0700 Subject: [PATCH 27/67] [CI] Wire untethered test files into CI jobs (#49340) Signed-off-by: Nick Hill Co-authored-by: Claude Opus 4.8 (1M context) --- .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/disaggregated.yaml | 16 +++++++++ .buildkite/test_areas/engine.yaml | 2 ++ .buildkite/test_areas/kernels.yaml | 36 +++++++++++++++++++ .buildkite/test_areas/misc.yaml | 10 ++++-- .buildkite/test_areas/models_basic.yaml | 3 +- .buildkite/test_areas/spec_decode.yaml | 2 ++ ..._fused_minimax_m3_qknorm_rope_kv_insert.py | 21 +++++++---- .../nixl_integration/run_edge_case_test.sh | 10 +++--- 9 files changed, 88 insertions(+), 13 deletions(-) diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index 927b5bd27f2..431ce07af4d 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -16,6 +16,7 @@ steps: commands: - pytest -v -s cuda/test_cuda_context.py - pytest -v -s cuda/test_platform_no_cuda_init.py + - pytest -v -s cuda/test_cuda_compatibility_path.py - label: Cudagraph device: h200_35gb diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index a3342e362ed..f1a89b39682 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -131,6 +131,22 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: NixlConnector PD edge case test (2 GPUs) + key: nixlconnector-pd-edge-cases-2-gpus + timeout_in_minutes: 40 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - tests/v1/kv_connector/nixl_integration/ + env: + PREFILL_GPU_ID: "0" + DECODE_GPU_ID: "1" + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_edge_case_test.sh + - label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus timeout_in_minutes: 25 diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index ed593c4aba2..ce4fc590eec 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -40,9 +40,11 @@ steps: source_file_dependencies: - vllm/v1/engine/ - tests/v1/engine/ + - tests/v1/test_tensor_ipc_queue.py commands: - pytest -v -s v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py + - pytest -v -s v1/test_tensor_ipc_queue.py mirror: amd: device: mi250_1 diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e1951685f60..938f8690551 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -61,9 +61,45 @@ steps: source_file_dependencies: - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu - vllm/models/deepseek_v4/common/ops/ + - vllm/models/deepseek_v4/nvidia/ - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + - tests/models/test_deepseek_v4_mega_moe.py commands: - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s models/test_deepseek_v4_mega_moe.py + +# Catch-all for test files at the tests/kernels root. This job collects +# the whole root so new files are wired by default. +# Files with dedicated jobs elsewhere in this file are excluded via --ignore +# (test_kda, test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in +# their own jobs / Kernels (B200)). +- label: Kernels Root Misc Test (B200) + key: kernels-root-misc-test-b200 + timeout_in_minutes: 45 + device: b200-k8s + source_file_dependencies: + - csrc/ + - vllm/ + - tests/kernels/ + commands: + - pytest -v -s kernels/ + --ignore=kernels/attention + --ignore=kernels/core + --ignore=kernels/helion + --ignore=kernels/ir + --ignore=kernels/mamba + --ignore=kernels/moe + --ignore=kernels/quantization + --ignore=kernels/test_concat_mla_q.py + --ignore=kernels/test_fused_qk_norm_rope_gate.py + --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py + --ignore=kernels/test_top_k_per_row.py + --ignore=kernels/test_kda.py + --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py + --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_shuffle_rows.py + # BROKEN on main, pending kernel fixes (B200): + # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) - label: Kernels Attention Test %N key: kernels-attention-test diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 1a53a92961f..763fbcbfec2 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -148,6 +148,7 @@ steps: - pytest -v -s -m 'cpu_test' v1/core - pytest -v -s v1/structured_output - pytest -v -s v1/test_serial_utils.py + - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics @@ -265,6 +266,7 @@ steps: - vllm/utils/ - vllm/v1/ - tests/v1/tracing + - tests/tracing/ commands: - "pip install \ 'opentelemetry-sdk>=1.26.0' \ @@ -272,6 +274,7 @@ steps: 'opentelemetry-exporter-otlp>=1.26.0' \ 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing + - pytest -v -s tracing mirror: amd: dind: false @@ -425,7 +428,7 @@ steps: - label: Batch Invariance (B200) key: batch-invariance-b200 - timeout_in_minutes: 35 + timeout_in_minutes: 45 device: b200-k8s source_file_dependencies: - vllm/v1/attention @@ -440,7 +443,10 @@ steps: - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py - + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: Acceptance Length Test (Large Models) # optional device: h200_35gb key: acceptance-length-test-large-models diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index a7c7aa9022d..24bfff5d756 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -82,7 +82,8 @@ steps: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/test_adapters.py - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index c63aaa18d8b..7cde7124cdc 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -90,8 +90,10 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ + - tests/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix" + - python3 spec_decode/test_custom_proposer.py mirror: amd: dind: false diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 626b06290e0..9c4a996438e 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -147,8 +147,11 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) # V is untouched. torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) @@ -255,8 +258,11 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps ).view(num_tokens, HEAD_DIM) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2) torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2) @@ -376,8 +382,11 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype): eps, ).view(num_tokens, kvsz) - torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2) + # The fused kernel keeps an fp32 intermediate across norm->rope, while the + # reference materializes bf16 after the norm (the unfused boundary), so + # rounding-boundary elements can differ by ~1 bf16 ulp. + torch.testing.assert_close(q_out, q_ref, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(k_out, k_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) torch.testing.assert_close(index_q_out, index_q_in, rtol=0, atol=0) torch.testing.assert_close(index_k_out, index_k_in, rtol=0, atol=0) diff --git a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh index 9d8e4df8c53..c3240ab5c17 100755 --- a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh @@ -3,8 +3,8 @@ set -xe # Parse command line arguments KV_BUFFER_DEVICE="cuda" # Default to cuda -PREFILL_GPU_ID=4 # Default GPU IDs -DECODE_GPU_ID=5 +PREFILL_GPU_ID="${PREFILL_GPU_ID:-4}" # Default GPU IDs +DECODE_GPU_ID="${DECODE_GPU_ID:-5}" while [[ $# -gt 0 ]]; do case $1 in --kv_buffer_device) @@ -70,6 +70,7 @@ run_tests_for_model() { --port $PREFILL_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -84,6 +85,7 @@ run_tests_for_model() { --port $DECODE_PORT \ --enforce-eager \ --gpu-memory-utilization 0.2 \ + --max-model-len 8192 \ --kv-transfer-config '$KV_CONFIG'" FULL_CMD="$BASE_CMD" @@ -98,7 +100,7 @@ run_tests_for_model() { # Build the command for the proxy server with all the hosts and ports PROXY_PORT=8192 - PROXY_CMD="python ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" + PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port $PROXY_PORT" PROXY_CMD+=" --prefiller-ports ${PREFILL_PORT}" PROXY_CMD+=" --decoder-ports ${DECODE_PORT}" # Start the proxy server @@ -110,7 +112,7 @@ run_tests_for_model() { # Run lm eval for this model echo "Running tests for $model_name" - PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py + PREFILL_PORT=$PREFILL_PORT DECODE_PORT=$DECODE_PORT PROXY_PORT=$PROXY_PORT python3 -m pytest -s -v "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_edge_cases.py # Clean up before running next model cleanup_instances From b6cbba8bc893c61e412a205533aafbee1ae6be31 Mon Sep 17 00:00:00 2001 From: oops-oom Date: Tue, 28 Jul 2026 22:24:02 +0800 Subject: [PATCH 28/67] [Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size (#48391) Signed-off-by: oops-oom <73481342@qq.com> Signed-off-by: oops-oom Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Shengqi Chen --- .buildkite/test_areas/misc.yaml | 14 +- csrc/libtorch_stable/layernorm_kernels.cu | 14 +- .../layernorm_quant_kernels.cu | 9 +- ...fused_layernorm_dynamic_per_token_quant.cu | 5 +- tests/v1/determinism/test_batch_invariance.py | 18 +++ .../test_rms_norm_batch_invariant.py | 150 +++++++++++++++++- 6 files changed, 189 insertions(+), 21 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 763fbcbfec2..caa56c21b37 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -398,7 +398,7 @@ steps: - label: Batch Invariance (A100) key: batch-invariance-a100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: a100 source_file_dependencies: - vllm/v1/attention @@ -408,11 +408,11 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: Batch Invariance (H100) key: batch-invariance-h100 - timeout_in_minutes: 40 + timeout_in_minutes: 60 device: h100 source_file_dependencies: - vllm/v1/attention @@ -423,8 +423,8 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - label: Batch Invariance (B200) key: batch-invariance-b200 @@ -439,8 +439,8 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA] - - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN] + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py - pytest -v -s v1/determinism/test_matmul_batch_invariant.py diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 878b44df936..7a1051d2c00 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -249,7 +249,9 @@ void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -325,8 +327,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] /* This kernel is memory-latency bound in many scenarios. When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency - hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + hiding on global mem ops. In batch-invariant mode the block size must + not depend on num_tokens, otherwise the same token would use a different + reduction width (and thus a different floating-point summation order) + across batches; lock it to 1024 to keep results bit-exact. */ + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -337,7 +344,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] auto res_ptr = reinterpret_cast(residual.data_ptr()); bool offsets_are_multiple_of_vector_width = hidden_size % vector_width == 0 && input_stride % vector_width == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); const bool has_weight = weight.has_value(); if (has_weight) { auto wt_ptr = reinterpret_cast(weight->data_ptr()); diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index f3bf8882e77..f43be531de0 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant( int num_tokens = input.numel() / hidden_size; // For large num_tokens, use smaller blocks to increase SM concurrency. - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 grid(num_tokens); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant( When num_tokens is large, a smaller block size allows for increased block occupancy on CUs and better latency hiding on global mem ops. */ - const int max_block_size = (num_tokens < 256) ? 1024 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant( auto wt_ptr = reinterpret_cast(weight.data_ptr()); bool ptrs_are_aligned = inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0; - bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 && !batch_invariant_launch) { LAUNCH_FUSED_ADD_RMS_NORM(8); diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 2152e64dc96..56dd4703872 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -2,6 +2,7 @@ #include "../../torch_utils.h" #include "../../dispatch_utils.h" +#include "../../../core/batch_invariant.hpp" #include "layernorm_utils.cuh" #include "quant_conversions.cuh" @@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch( auto num_tokens = input.numel() / hidden_size; dim3 grid(num_tokens); - const int max_block_size = (num_tokens <= 256) ? 512 : 256; + const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); + const int max_block_size = + batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256); dim3 block(std::min(hidden_size, max_block_size)); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b2706ed89b7..37fd5cba6a5 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -27,8 +27,10 @@ from vllm.platforms import current_platform "backend", BACKENDS, ) +@pytest.mark.parametrize("rms_norm_impl", ["default", "vllm_c"]) def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( backend, + rms_norm_impl, ): """ Ensures that the same request (the 'needle' prompt) yields identical output @@ -60,6 +62,16 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( random.seed(seed) attention_config = {"backend": backend} + # Force the C++ RMSNorm implementation so we actually exercise the + # num_tokens-dependent block-size branches. + kernel_config = None + if rms_norm_impl == "vllm_c": + kernel_config = { + "ir_op_priority": { + "rms_norm": ["vllm_c"], + "fused_add_rms_norm": ["vllm_c"], + } + } # Allow overrides from environment (useful for CI tuning) # "facebook/opt-125m" is too small, doesn't reliably test determinism model = TEST_MODEL @@ -96,6 +108,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( gpu_memory_utilization=gpu_mem_util, max_model_len=max_model_len, attention_config=attention_config, + kernel_config=kernel_config, ) # Baseline generation for the needle prompt alone. @@ -923,11 +936,15 @@ def LLM_with_max_seqs( gpu_memory_utilization: float, max_model_len: int, attention_config: dict | None = None, + kernel_config: dict | None = None, ) -> LLM: """ Helper to construct an LLM with a specific max_num_seqs (batch-size limit) using the high-level v1 LLM API, while constraining memory usage. """ + extra_kwargs: dict = {} + if kernel_config is not None: + extra_kwargs["kernel_config"] = kernel_config return LLM( model=model, max_num_seqs=max_num_seqs, @@ -939,4 +956,5 @@ def LLM_with_max_seqs( attention_config=attention_config, # Enable for MOE models # enable_expert_parallel=True, + **extra_kwargs, ) diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index dfd08351277..232a43b1f98 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -28,16 +28,18 @@ def _rms_norm_reference( @skip_if_not_cuda -@pytest.mark.parametrize("batch_size", [1, 4, 16, 64]) +@pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6, 1e-5]) +@pytest.mark.parametrize("seed", list(range(4))) def test_rms_norm_batch_invariant_vs_reference( default_vllm_config, batch_size: int, hidden_size: int, dtype: torch.dtype, eps: float, + seed: int, ): """ Compare batch-invariant Triton RMS norm against a PyTorch reference. @@ -48,7 +50,7 @@ def test_rms_norm_batch_invariant_vs_reference( device = torch.device(DEVICE_TYPE) # Create test input and weight - torch.manual_seed(42) + torch.manual_seed(seed) input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) @@ -71,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_reference( atol=atol, msg=f"RMS norm mismatch for batch_size={batch_size}, " f"hidden_size={hidden_size}, " - f"dtype={dtype}, eps={eps}", + f"dtype={dtype}, eps={eps}, seed={seed}", ) @@ -79,17 +81,21 @@ def test_rms_norm_batch_invariant_vs_reference( @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("n_extra", [3, 299]) +@pytest.mark.parametrize("seed", list(range(16))) def test_fused_add_rms_norm_batch_invariant_residual_path( hidden_size: int, dtype: torch.dtype, eps: float, + n_extra: int, + seed: int, ): """ Test the batch-invariant fused residual-add + RMSNorm helper directly. """ device = torch.device(DEVICE_TYPE) - torch.manual_seed(42) + torch.manual_seed(seed) x_single = torch.randn(1, hidden_size, dtype=dtype, device=device) residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device) weight = torch.randn(hidden_size, dtype=dtype, device=device) @@ -97,14 +103,14 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( x_batch = torch.cat( [ x_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) residual_batch = torch.cat( [ residual_single, - torch.randn(3, hidden_size, dtype=dtype, device=device), + torch.randn(n_extra, hidden_size, dtype=dtype, device=device), ], dim=0, ) @@ -168,6 +174,138 @@ def test_fused_add_rms_norm_batch_invariant_residual_path( ) +FP8_DTYPE = current_platform.fp8_dtype() + +# The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block +# 256, while the small launch (255 rows) stays under the threshold and keeps the +# larger block (1024, or 512 for per-block quant). Under the pin the two launches +# use the same block, so the shared first 255 rows must match bit-for-bit; 255 is +# the most rows a single small launch can hold (< 256, and <= 256 for per-block). +_LARGE_TOKENS = 300 +_SMALL_TOKENS = 255 + + +def _assert_rows_bit_identical(small, large, msg): + if small.dtype == FP8_DTYPE: + assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg + else: + torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_batch_invariant_nonresidual_kernel( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm`` (no residual) must be batch invariant across the block + threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]`` + (default priority is ``native``/inductor codegen when compiling). + """ + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + + def rms_norm(x): + out = torch.empty_like(x) + ops.rms_norm(out, x, weight, 1e-6) + return out + + large = rms_norm(rows.clone()) + small = rms_norm(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "rms_norm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +@pytest.mark.parametrize("add_residual", [False, True]) +def test_rms_norm_static_fp8_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool +): + """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across + the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with + ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8 + path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them). + """ + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + residual = ( + torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + if add_residual + else None + ) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + + def quant(x, res): + out = torch.empty_like(x, dtype=FP8_DTYPE) + if add_residual: + torch.ops._C.fused_add_rms_norm_static_fp8_quant( + out, x, res, weight, quant_scale, 1e-6 + ) + else: + torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6) + return out + + large = quant(rows.clone(), residual.clone() if residual is not None else None) + small = quant( + rows[:_SMALL_TOKENS].clone(), + residual[:_SMALL_TOKENS].clone() if residual is not None else None, + ) + _assert_rows_bit_identical( + small, + large[:_SMALL_TOKENS], + "static-fp8-quant RMSNorm output depends on num_tokens (block size)", + ) + + +@skip_if_not_cuda +@pytest.mark.parametrize("hidden_size", [512, 4096]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("seed", list(range(4))) +def test_rms_norm_per_block_quant_batch_invariant( + hidden_size: int, dtype: torch.dtype, seed: int +): + """C++ ``rms_norm_per_block_quant`` must be batch invariant across the + block threshold (compiled fp8 block-quant path; block pinned to 512).""" + import vllm._custom_ops as ops + + device = torch.device(DEVICE_TYPE) + torch.manual_seed(seed) + rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device) + weight = torch.randn(hidden_size, dtype=dtype, device=device) + group_size = [1, 128] + + def per_block_quant(x): + return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size) + + out_large, scale_large = per_block_quant(rows.clone()) + out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone()) + _assert_rows_bit_identical( + out_small, + out_large[:_SMALL_TOKENS], + "rms_norm_per_block_quant output depends on num_tokens (block size)", + ) + torch.testing.assert_close( + scale_small, + scale_large[:_SMALL_TOKENS], + rtol=0.0, + atol=0.0, + msg="rms_norm_per_block_quant scales depend on num_tokens (block size)", + ) + + @skip_if_not_cuda @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) From 1e81853afc8701cb50b649e94a23c977da8c1ed0 Mon Sep 17 00:00:00 2001 From: Jonguk Cheong Date: Tue, 28 Jul 2026 23:41:55 +0900 Subject: [PATCH 29/67] [Bugfix][KV Offload] Keep Mamba block span unscaled under DCP (#49964) Signed-off-by: Jonguk Cheong Co-authored-by: OpenAI Codex --- .../unit/offloading_connector/test_config.py | 56 +++++++++++++++++++ .../kv_connector/v1/offloading/config.py | 12 +++- 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index fc426ff318d..5a66b463e30 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -8,10 +8,15 @@ from unittest.mock import MagicMock, patch import pytest import torch +from tests.v1.kv_connector.unit.offloading_connector.utils import MockOffloadingSpec from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( build_offloading_config, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + SchedulerOffloadConfig, + is_store_reachable_swa_chunk, +) from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -179,6 +184,25 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig: ) +def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["full_layer"], _full_attention_spec()), + KVCacheGroupSpec( + ["mamba_layer"], + MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + + def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: config = _make_vllm_config() kv_cache_config = KVCacheConfig( @@ -267,6 +291,38 @@ def test_prefill_context_parallelism_does_not_scale_group_blocks(): assert offloading_config.cache.blocks_per_chunk == 4 +def test_dcp_scales_attention_but_not_mamba_group_blocks(): + config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2) + config.speculative_config = None + + offloading_config = build_offloading_config( + config, _make_mamba_hybrid_kv_cache_config() + ) + + assert tuple(group.tokens_per_block for group in offloading_config.groups) == ( + 32, + 16, + ) + scheduler_config = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), + config, + _make_mamba_hybrid_kv_cache_config(), + ) + mamba_group = scheduler_config.kv_group_configs[1] + assert mamba_group.alignment_chunk_count == 2 + assert [ + chunk_idx + for chunk_idx in range(4) + if is_store_reachable_swa_chunk( + chunk_idx, + 4, + mamba_group.alignment_chunk_count, + mamba_group.sliding_window_size_in_chunks, + mamba_group.is_eagle_group, + ) + ] == [1, 3] + + def test_preserves_data_parallel_index(): config = _make_vllm_config() config.parallel_config.data_parallel_index = 2 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b86ebf96bb6..b9837bcb4b0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -5,7 +5,11 @@ from typing import TYPE_CHECKING from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, + MLAAttentionSpec, +) from vllm.v1.kv_offload.config import ( OffloadingCacheConfig, OffloadingConfig, @@ -40,7 +44,11 @@ def build_offloading_config( OffloadingGroupConfig( tokens_per_block=( group.kv_cache_spec.block_size - * parallel_config.decode_context_parallel_size + * ( + parallel_config.decode_context_parallel_size + if isinstance(group.kv_cache_spec, AttentionSpec) + else 1 + ) ), layer_names=tuple(group.layer_names), ) From 0d0504b54c73119ed643c80b4ed56ef3cf80e209 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 07:58:02 -0700 Subject: [PATCH 30/67] [Core] Warm up runner-owned Triton kernels before the first request (#49903) --- tests/v1/worker/test_kv_block_zeroer.py | 66 +++++++++- vllm/model_executor/warmup/kernel_warmup.py | 16 ++- .../warmup/qwen_triton_warmup.py | 114 ----------------- .../warmup/v1_block_table_warmup.py | 42 ++---- vllm/v1/worker/gpu/warmup.py | 121 ++++++++++++------ vllm/v1/worker/mamba_utils.py | 6 +- vllm/v1/worker/utils.py | 7 +- 7 files changed, 183 insertions(+), 189 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index b212e3ae17b..17aa1bf38d4 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.worker.utils import KVBlockZeroer +from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -86,3 +86,67 @@ def test_non_uniform_page_sizes(): assert torch.all(storage[1] == 0) assert torch.all(storage[2] == 0) assert torch.all(storage[3] == 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_compiles_every_n_blocks_specialization(): + """After warmup, no launch should trigger a first-request JIT compile. + + ``n_blocks`` is ``do_not_specialize``, so a single warmup launch must + cover every block count. + """ + device = torch.device("cuda") + num_blocks = 64 + page_size_el = 4 + storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, # max_chunks + page_size_el, # blk_size + 1, # n_segs + ) + + def compiled_variants() -> set: + return { + key + for caches in _zero_kv_blocks_kernel.device_caches.values() + for key in caches[0] + } + + zeroer.warmup(num_blocks) + torch.accelerator.synchronize() + warmed = compiled_variants() + assert warmed + + for n_blocks in (1, 2, 3, 16, 32): + zeroer.zero_block_ids(list(range(n_blocks))) + torch.accelerator.synchronize() + + assert compiled_variants() == warmed + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_warmup_respects_available_block_count(): + """An empty KV cache must not be warmed with out-of-range block IDs.""" + device = torch.device("cuda") + page_size_el = 4 + storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device) + + zeroer = KVBlockZeroer.__new__(KVBlockZeroer) + zeroer.device = device + zeroer._meta = ( + torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device), + torch.tensor([page_size_el], dtype=torch.int64, device=device), + 1, + page_size_el, + 1, + ) + + zeroer.warmup(0) + torch.accelerator.synchronize() + + assert torch.all(storage == 1) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b2c989b295c..afe47f39ced 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -98,12 +98,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): minimax_m3_msa_warmup, ) - # Pooling models do not use the generation slot-mapping path. - if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels( - getattr(worker.model_runner, "device", torch.device("cuda")), - worker.scheduler_config.max_num_batched_tokens, - ) + if not worker.use_v2_model_runner: + # Pooling models do not use the generation slot-mapping path. + if not worker.model_runner.is_pooling_model: + warm_v1_block_table_kernels(worker.model_runner) + # The KV-block zeroing kernel is driven by the scheduler's + # `new_block_ids_to_zero`, so no dummy run ever reaches it. + zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) + if zeroer is not None: + zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index 8cbfa539b7e..06a016a5a06 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -24,24 +24,10 @@ _QWEN_MODEL_TYPES = frozenset( } ) -_ZERO_KV_N_BLOCKS = (1, 2) - -_SLOT_MAPPING_KV_BLOCK_SIZE = 16 -_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3) - # Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L. _FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16) -@dataclass(frozen=True) -class _ZeroKvWarmupConfig: - seg_page_sizes: torch.Tensor - max_chunks: int - block_size: int - n_segs: int - - @dataclass(frozen=True) class _QwenGDNWarmupConfig: h: int @@ -148,96 +134,6 @@ def _qwen_gdn_warmup_config( return None -def _get_kv_block_zeroer(runner: object) -> object | None: - zeroer = getattr(runner, "kv_block_zeroer", None) - if zeroer is None: - zeroer = getattr(runner, "_kv_block_zeroer", None) - return zeroer - - -def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None: - zeroer = _get_kv_block_zeroer(runner) - meta = getattr(zeroer, "_meta", None) - if meta is None: - return None - - _, seg_page_sizes, max_chunks, block_size, n_segs = meta - return _ZeroKvWarmupConfig( - seg_page_sizes=seg_page_sizes, - max_chunks=int(max_chunks), - block_size=int(block_size), - n_segs=int(n_segs), - ) - - -def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool: - zeroer = _get_kv_block_zeroer(runner) - zero_block_ids = getattr(zeroer, "zero_block_ids", None) - if not callable(zero_block_ids): - return False - - for n_blocks in _ZERO_KV_N_BLOCKS: - zero_block_ids(list(range(n_blocks))) - return True - - -def _warm_zero_kv_blocks_kernel( - device: torch.device, config: _ZeroKvWarmupConfig -) -> None: - from vllm.v1.worker.utils import _zero_kv_blocks_kernel - - max_n_blocks = max(_ZERO_KV_N_BLOCKS) - max_page_size = int(config.seg_page_sizes.max().item()) - scratch = torch.empty( - max_n_blocks * max_page_size, - dtype=torch.int32, - device=device, - ) - seg_addrs = torch.tensor( - [scratch.data_ptr()] * config.n_segs, - dtype=torch.uint64, - device=device, - ) - - for n_blocks in _ZERO_KV_N_BLOCKS: - block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device) - grid = (n_blocks * config.n_segs * config.max_chunks,) - _zero_kv_blocks_kernel[grid]( - seg_addrs, - config.seg_page_sizes, - block_ids, - n_blocks, - N_SEGS=config.n_segs, - MAX_CHUNKS=config.max_chunks, - BLOCK_SIZE=config.block_size, - ) - - -def _warm_compute_slot_mapping_kernel(device: torch.device) -> None: - from vllm.v1.worker.block_table import BlockTable - - # num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny. - num_tokens = 1 - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - - for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES: - # Use BlockTable so the JIT key matches the production slot-mapping call. - block_table = BlockTable( - block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - max_num_reqs=1, - max_num_blocks_per_req=block_table_stride, - max_num_batched_tokens=num_tokens, - pin_memory=False, - device=device, - kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE, - cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE, - ) - block_table.add_row(list(range(block_table_stride)), 0) - block_table.commit_block_table(num_reqs=1) - block_table.compute_slot_mapping(1, query_start_loc, positions) - - def _warm_causal_conv1d_fwd_kernel( device: torch.device, config: _QwenGDNWarmupConfig ) -> None: @@ -373,16 +269,6 @@ def qwen_triton_warmup( device = getattr(runner, "device", torch.device("cuda")) logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) - zero_config = _zero_kv_warmup_config(runner) - warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner) - if zero_config is not None: - _warm_zero_kv_blocks_kernel(device, zero_config) - elif not warmed_zeroer: - logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") - - _warm_compute_slot_mapping_kernel(device) - _synchronize_device(device) - compilation_config = getattr(runner, "compilation_config", None) static_forward_context = getattr(compilation_config, "static_forward_context", None) gdn_config = _qwen_gdn_warmup_config(static_forward_context) diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py index 8d2328432eb..d49e1ba7cc8 100644 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ b/vllm/model_executor/warmup/v1_block_table_warmup.py @@ -2,42 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm up v1 block-table Triton kernels.""" +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + _SLOT_MAPPING_WARMUP_TOKENS = 8 -_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16) -_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1 -def warm_v1_block_table_kernels( - device: torch.device, - max_tokens: int, -) -> None: - from vllm.v1.worker.block_table import BlockTable +def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: + """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens)) + device = runner.device + block_table = runner.input_batch.block_table + num_tokens = min( + _SLOT_MAPPING_WARMUP_TOKENS, + runner.scheduler_config.max_num_batched_tokens, + ) if num_tokens <= 0: return query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES: - max_num_blocks_per_req = max( - 1, (max(num_tokens, max_tokens) + block_size - 1) // block_size - ) - max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16 - block_table = BlockTable( - block_size=block_size, - max_num_reqs=1, - max_num_blocks_per_req=max_num_blocks_per_req, - max_num_batched_tokens=max(num_tokens, max_tokens), - pin_memory=False, - device=device, - kernel_block_size=block_size, - cp_kv_cache_interleave_size=( - _SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE - ), - ) - block_table.add_row(list(range(max_num_blocks_per_req)), 0) - block_table.commit_block_table(1) - block_table.compute_slot_mapping(1, query_start_loc, positions) + block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 785188dc3c3..260ca7a5f17 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -157,13 +157,10 @@ def warmup_kernels( worker_execute_model: Callable[[SchedulerOutput], Any], worker_sample_tokens: Callable[[GrammarOutput | None], Any], ) -> None: - """Run two execute_model + sample_tokens iterations to JIT compile - triton kernels. We must call the provided worker's execute_model for - pipeline parallel coordination. + """Run scheduler-realistic prefill and decode steps to JIT compile kernels. - The first iteration simulates a prefill with requests of - decode_query_len + 1 prompt tokens each. The second iteration simulates - a decode step with all requests generating decode_query_len tokens. + We must call the provided worker's execute_model for pipeline parallel + coordination. """ num_spec_steps = model_runner.num_speculative_steps decode_query_len = model_runner.decode_query_len @@ -172,8 +169,13 @@ def warmup_kernels( # a uniform decode batch. prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates decode_query_len tokens. - decode_len = prompt_len + decode_query_len + # Upper bound on the decode steps built in `decode_steps` below. + num_decode_steps = 1 + if not model_runner.is_pooling_model: + num_decode_steps = 5 if num_spec_steps > 0 else 3 + # Size the block allocation for the worst case: every request advancing + # decode_query_len tokens on every decode step. + decode_len = prompt_len + num_decode_steps * decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) @@ -208,9 +210,6 @@ def warmup_kernels( kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups] prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs] decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs] - decode_block_deltas = [ - d - p for d, p in zip(decode_block_counts, prefill_block_counts) - ] max_blocks_per_req = sum(decode_block_counts) num_reqs = min( @@ -243,6 +242,11 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) + # The KV-block zeroing kernel is driven by the scheduler's + # new_block_ids_to_zero, so none of the steps below reach it. + if model_runner.kv_block_zeroer is not None: + model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks) + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( @@ -287,33 +291,78 @@ def warmup_kernels( worker_sample_tokens(grammar_output) - # Step 2: Decode all requests with decode_query_len tokens each. - cached_req_data = CachedRequestData.make_empty() - cached_req_data.req_ids = list(req_ids) - cached_req_data.num_computed_tokens = [prompt_len] * num_reqs - cached_req_data.num_output_tokens = [1] * num_reqs - new_block = any(decode_block_deltas) - cached_req_data.new_block_ids = [ - tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None - for _ in range(num_reqs) + # Per-request state carried across the decode steps. + req_computed = [prompt_len] * num_reqs + req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)] + + def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None: + """Decode `indices`, spec-decoding the ones flagged in `spec_flags`.""" + cached_req_data = CachedRequestData.make_empty() + cached_req_data.req_ids = [req_ids[i] for i in indices] + cached_req_data.num_computed_tokens = [req_computed[i] for i in indices] + cached_req_data.num_output_tokens = [1] * len(indices) + cached_req_data.new_block_ids = [] + + step_num_scheduled_tokens: dict[str, int] = {} + step_spec_tokens: dict[str, list[int]] = {} + for i, use_spec in zip(indices, spec_flags): + num_tokens = decode_query_len if use_spec else 1 + after = req_computed[i] + num_tokens + deltas = [ + _warmup_block_count(after, spec) - held + for spec, held in zip(kv_cache_specs, req_blocks[i]) + ] + cached_req_data.new_block_ids.append( + tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None + ) + req_blocks[i] = [ + held + delta for held, delta in zip(req_blocks[i], deltas) + ] + step_num_scheduled_tokens[req_ids[i]] = num_tokens + if use_spec: + step_spec_tokens[req_ids[i]] = [0] * num_spec_steps + + decode_output = SchedulerOutput.make_empty() + decode_output.scheduled_cached_reqs = cached_req_data + decode_output.num_scheduled_tokens = step_num_scheduled_tokens + decode_output.scheduled_spec_decode_tokens = step_spec_tokens + decode_output.total_num_scheduled_tokens = sum( + step_num_scheduled_tokens.values() + ) + decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + worker_execute_model(decode_output) + worker_sample_tokens(None) + + for i, use_spec in zip(indices, spec_flags): + req_computed[i] += decode_query_len if use_spec else 1 + + all_indices = list(range(num_reqs)) + use_spec_decode = num_spec_steps > 0 + + # Decode steps to warm, as (request indices, per-request spec flag). + # Under spec decoding the scheduler drops requests the drafter proposed + # nothing for, so warm each batch shape with and without draft tokens. + decode_steps: list[tuple[list[int], list[bool]]] = [ + (all_indices, [use_spec_decode] * num_reqs), ] + if num_reqs >= 2: + # Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode + # as a prefill and split the batch into spec/non-spec token indices. + decode_steps.append(([0, 1], [use_spec_decode, False])) + if use_spec_decode: + # Exercise the model paths that split a batch by whether each + # request received draft tokens. + decode_steps.append(([0, 1], [False, False])) + if num_reqs > 1: + decode_steps.append(([0], [use_spec_decode])) + if use_spec_decode: + decode_steps.append(([0], [False])) + elif use_spec_decode: + decode_steps.append(([0], [False])) - decode_output = SchedulerOutput.make_empty() - decode_output.scheduled_cached_reqs = cached_req_data - decode_output.num_scheduled_tokens = { - req_id: decode_query_len for req_id in req_ids - } - if num_spec_steps > 0: - decode_output.scheduled_spec_decode_tokens = { - req_id: [0] * num_spec_steps for req_id in req_ids - } - decode_output.total_num_scheduled_tokens = sum( - decode_output.num_scheduled_tokens.values() - ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - - worker_execute_model(decode_output) - worker_sample_tokens(None) + for step_indices, step_spec_flags in decode_steps: + _run_decode_step(step_indices, step_spec_flags) # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7f611a2eeff..36b72efe35b 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -150,7 +150,7 @@ def _copy_mamba_state_block( tl.store(tail_dst + tail_off, tail_data, mask=tail_mask) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def postprocess_mamba_fused_kernel( # Decision inputs (per-request) num_accepted_tokens_ptr, @@ -280,7 +280,7 @@ def postprocess_mamba_fused_kernel( ) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def preprocess_mamba_align_fused_kernel( idx_mapping_ptr, state_idx_ptr, @@ -326,7 +326,7 @@ def preprocess_mamba_align_fused_kernel( tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) -@triton.jit +@triton.jit(do_not_specialize=["num_reqs"]) def precopy_mamba_align_fused_kernel( # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by # the V2 fused align preprocess kernel for the current step: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d00974046fc..4afb941c12a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) -@triton.jit +@triton.jit(do_not_specialize=["n_blocks"]) def _zero_kv_blocks_kernel( seg_addrs_ptr, seg_page_sizes_ptr, @@ -206,6 +206,11 @@ class KVBlockZeroer: BLOCK_SIZE=blk_size, ) + def warmup(self, num_kv_blocks: int) -> None: + """JIT-compile the zeroing kernel before the first real request.""" + if num_kv_blocks > 0: + self.zero_block_ids([0]) + @dataclass class AttentionGroup: From 30217b0e809ef045b3b4aa7c011d6c6d2b0b0362 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:01:55 +0300 Subject: [PATCH 31/67] [Bugfix][KV Offload][P2P] Scope serve state to fetch rounds (#49877) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis Co-authored-by: Itay Etelis --- .../v1/kv_offload/tiering/p2p/test_manager.py | 13 +- .../kv_offload/tiering/p2p/test_sessions.py | 116 ++++- .../kv_offload/tiering/p2p/session/client.py | 293 +++++------ .../tiering/p2p/session/protocol.py | 48 +- .../kv_offload/tiering/p2p/session/server.py | 474 ++++++++++-------- .../kv_offload/tiering/p2p/session/session.py | 16 +- 6 files changed, 570 insertions(+), 390 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index cd6e8f382a5..c8a9d7f3ea3 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -1612,16 +1612,17 @@ class TestBindHostPortDefaults: monkeypatch.setattr( manager_module, "NixlTransport", - lambda agent_name, *a, **k: calls.update(nixl_name=agent_name) - or SimpleNamespace(), + lambda agent_name, *a, **k: ( + calls.update(nixl_name=agent_name) or SimpleNamespace() + ), ) monkeypatch.setattr( manager_module, "ZmqTransport", - lambda local_id, host, port, *a, **k: calls.update( - zmq_id=local_id, zmq_host=host, zmq_port=port - ) - or SimpleNamespace(), + lambda local_id, host, port, *a, **k: ( + calls.update(zmq_id=local_id, zmq_host=host, zmq_port=port) + or SimpleNamespace() + ), ) spec = SimpleNamespace( blocks_per_chunk=1, diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index cab4ce6705a..20e4bad2186 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -49,6 +49,7 @@ from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( from vllm.v1.kv_offload.tiering.p2p.session.server import ( _CANCEL_DRAIN_TIMEOUT_S, _InflightXfer, + _OutboundRequestState, ) from vllm.v1.kv_offload.tiering.p2p.session.session import ( _MAX_CONSECUTIVE_DISPATCH_ERRORS, @@ -315,10 +316,26 @@ def _activate( # either a missing entry or a None field. These helpers paper over that. +def _client_load(session: P2PSession, kv_request_id: str): + """The single in-flight load of a kv_request_id (loads are per-round).""" + loads = session._client._requests[kv_request_id].loads + assert len(loads) == 1 + return next(iter(loads.values())) + + def _srv_outbound(session: P2PSession, kv_request_id: str): - """Outbound serve state for a kv_request_id, or None (idle / GC'd).""" + """Serve-side round for a kv_request_id, or None (idle / GC'd). + + Rounds are keyed by wire round_seq; surfaces the demanded round when + a fetch has bound one, else any parked supply round. + """ st = session._server._requests.get(kv_request_id) - return st.outbound if st is not None else None + if st is None or not st.outbound: + return None + for rnd in st.outbound.values(): + if rnd.demand_received: + return rnd + return next(iter(st.outbound.values())) def _srv_lookups(session: P2PSession) -> list: @@ -330,8 +347,10 @@ def _srv_lookups(session: P2PSession) -> list: def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: """Pending-abort start time for a kv_request_id, or None.""" - st = session._server._requests.get(kv_request_id) - return st.abort_started_at if st is not None else None + for (kv, _), started in session._server._pending_aborts.items(): + if kv == kv_request_id: + return started + return None def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: @@ -479,6 +498,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -495,6 +515,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: False, } @@ -538,6 +559,7 @@ class TestClientFlows: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -552,7 +574,7 @@ class TestClientFlows: session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) - session._client._requests["req-1"].load.submitted_at = time.monotonic() - 60.0 + _client_load(session, "req-1").submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE @@ -569,7 +591,7 @@ class TestClientFlows: job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. - session._client._requests["req-7"].load.submitted_at = ( + _client_load(session, "req-7").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -579,11 +601,11 @@ class TestClientFlows: and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) - assert session._client._requests["req-7"].load.aborted_at is not None + assert _client_load(session, "req-7").aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. - session._client._requests["req-7"].load.aborted_at = ( + _client_load(session, "req-7").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads @@ -599,17 +621,18 @@ class TestClientFlows: session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) - session._client._requests["req-8"].load.submitted_at = ( + _client_load(session, "req-8").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() - assert session._client._requests["req-8"].load.aborted_at is not None + assert _client_load(session, "req-8").aborted_at is not None # Peer acks the abort. conn.enqueue( { TYPE_KEY: AbortAckMsg.TYPE, + AbortAckMsg.ROUND_SEQ: 0, AbortAckMsg.KV_REQUEST_ID: "req-8", } ) @@ -915,6 +938,7 @@ class TestLookupFlow: conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: "req-1", LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], } @@ -951,6 +975,7 @@ def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, + LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: kv_request_id, LookupMsg.KEYS: list(keys), } @@ -1201,6 +1226,7 @@ class TestServerLookupHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [], FetchMsg.BLOCK_INDEXES: [], @@ -1240,6 +1266,7 @@ class TestServerLookupHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"hA", b"hB"], FetchMsg.BLOCK_INDEXES: [20, 21], @@ -1277,6 +1304,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1295,6 +1323,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1313,6 +1342,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1331,6 +1361,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1349,13 +1380,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1376,13 +1413,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1409,20 +1452,26 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. - session._server._requests["req-1"].abort_started_at = ( + session._server._pending_aborts[("req-1", 0)] = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the @@ -1445,13 +1494,19 @@ class TestServerFlows: tid = 42 session._server._inflight_add( tid, - _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}), + _InflightXfer( + kv_request_id="req-1", + block_count=1, + job_ids={1}, + round=_OutboundRequestState(inflight=1), + ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1463,6 +1518,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) @@ -1497,6 +1553,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1529,6 +1586,7 @@ class TestServerFlows: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1571,6 +1629,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], @@ -1596,6 +1655,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1631,6 +1691,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1665,6 +1726,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1686,6 +1748,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-2", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1719,6 +1782,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], @@ -1753,6 +1817,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1790,6 +1855,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], @@ -1824,6 +1890,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], @@ -1885,6 +1952,7 @@ class TestFinishRequestServerSide: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], @@ -1951,6 +2019,7 @@ class TestBidirectional: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-srv", FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], @@ -1981,6 +2050,7 @@ class TestBidirectional: conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-cli", TransferDoneMsg.SUCCESS: True, } @@ -2168,6 +2238,7 @@ class TestAdversarial: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], @@ -2216,6 +2287,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], @@ -2246,6 +2318,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2269,6 +2342,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2296,6 +2370,7 @@ class TestDispatchErrorHandling: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], @@ -2341,6 +2416,7 @@ class TestInflightPerReqInvariant: conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: kv_id, FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, @@ -2394,7 +2470,12 @@ class TestInflightPerReqInvariant: tid = kv_id_idx * 10 + j session._server._inflight_add( tid, - _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}), + _InflightXfer( + kv_request_id=kv_id, + block_count=1, + job_ids={tid}, + round=_OutboundRequestState(inflight=1), + ), ) assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") @@ -2488,6 +2569,7 @@ class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, + FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], @@ -2513,6 +2595,7 @@ class TestTransferDoneMsgValidation: def test_valid_message_passes(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } @@ -2521,6 +2604,7 @@ class TestTransferDoneMsgValidation: def test_success_wrong_type(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: 1, } diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index ed91f7f9b00..07a4620780e 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -11,7 +11,6 @@ callback injected by the coordinator (which gates on ConnectAck). from __future__ import annotations -import enum import time from collections.abc import Callable, Sequence from dataclasses import dataclass, field @@ -35,26 +34,12 @@ _LOAD_TIMEOUT_S = 30.0 _ABORT_ACK_TIMEOUT_S = 10.0 -class ClientPhase(enum.Enum): - """Lifecycle of a request's client-side lookup/fetch signalling. - - Advances monotonically. Only ``finish`` reads it, to decide - whether a terminal empty FetchMsg is owed to release the peer's - lookup state (owed only from ``PROBING``: a LookupMsg went out but no - FetchMsg has since closed the peer's lookup phase). - """ - - REGISTERED = enum.auto() # keys in probes/unsent, nothing sent yet - PROBING = enum.auto() # LookupMsg flushed, awaiting responses - FETCH_SENT = enum.auto() # FetchMsg sent (real or terminal empty) - - @dataclass class _InboundLoadState: """Client-role state for a single in-flight load request. - Lives on ``_ClientRequestState.load`` for the duration of a fetch; - the owning kv_request_id is the dict key, so it isn't stored here. + Lives in ``_ClientRequestState.loads`` keyed by round_seq for the + duration of a fetch; the owning kv_request_id is the outer dict key. """ job_id: int # opaque ID assigned by the manager to this load request @@ -68,7 +53,7 @@ class _ClientRequestState: One entry per kv_request_id we're driving. Lookup-phase fields are used only by symmetric P2P (``do_p2p_fetch``); PD-only loads leave - ``probes``/``unsent`` empty and drive just ``phase`` and ``load``. An + ``probes``/``unsent`` empty and drive just ``phase`` and ``loads``. An entry is dropped once every field is idle — see ``ClientRole._maybe_prune``. """ @@ -82,11 +67,23 @@ class _ClientRequestState: # OffloadKeys registered but not yet flushed onto the wire. Drained and # cleared by the next flush_pending_lookups. unsent: list[OffloadKey] = field(default_factory=list) + # Current lookup round. LookupMsgs carry it, each fetch closes it and + # advances it, so every round's supply/demand/completion is isolated + # on the wire. PD clients never probe and stay on round 0. + round_seq: int = 0 + # This id ran the symmetric lookup phase (register_lookup); a fetch + # with keys then requires every key to be a confirmed probe. PD + # loads never probe. + probed: bool = False - # Monotonic lookup/fetch signalling phase; see ``ClientPhase``. - phase: ClientPhase = ClientPhase.REGISTERED - # Set while a fetch is in flight; cleared on completion/abort/timeout. - load: _InboundLoadState | None = None + # The peer holds lookup state no FetchMsg has closed: a LookupMsg + # was flushed since the last fetch. finish owes a terminal empty + # FetchMsg while set, so the peer releases parked supply. + peer_lookup_open: bool = False + # In-flight loads keyed by the round their fetch carried. The + # scheduler submits loads incrementally as chunks resolve, so several + # can be in flight at once; TransferDone/AbortAck match by round. + loads: dict[int, _InboundLoadState] = field(default_factory=dict) class LoadResult(NamedTuple): @@ -133,11 +130,10 @@ class ClientRole: # _serve_pending. Populated by register_lookup, drained by # flush_pending_lookups, and discarded on finish/close. self._flush_pending: set[str] = set() - # kv_request_ids with a fetch in flight (``st.load is not None``) — - # the work-list collect_results walks for timeouts, and the - # has_active_loads predicate, instead of scanning every request. - # Kept in exact sync with ``st.load``: armed in request_blocks, - # discarded wherever load is cleared, and cleared on close. + # kv_request_ids with at least one fetch in flight — the work-list + # collect_results walks for timeouts, and the has_active_loads + # predicate, instead of scanning every request. Kept in exact sync + # with ``st.loads``. self._active_loads: set[str] = set() self._completed_loads: list[LoadResult] = [] @@ -156,17 +152,21 @@ class ClientRole: def _maybe_prune(self, kv_request_id: str) -> None: """Drop the entry once it holds no live load or lookup state. - The sticky ``phase`` is only read by ``finish``. A probe - clears when its fetch is issued (``request_blocks``) or when the - request finishes (``finish``/``close``); in the former case - ``load`` is set and keeps the entry alive, in the latter the phase - is no longer needed — so dropping on emptiness never loses a phase - still in use. + ``peer_lookup_open`` is only read by ``finish``, and every path + that clears the last probe (fetch / finish / close) also settles + it, so dropping on emptiness never loses a flag still in use. """ st = self._requests.get(kv_request_id) - if st is not None and st.load is None and not st.probes and not st.unsent: + if st is not None and not st.loads and not st.probes and not st.unsent: del self._requests[kv_request_id] + def _on_load_terminal(self, kv_request_id: str, st: _ClientRequestState) -> None: + """Wind down id-level state once no load remains in flight.""" + if st.loads: + return + self._active_loads.discard(kv_request_id) + self._maybe_prune(kv_request_id) + @property def has_active_loads(self) -> bool: """True if any kv_request_id has a fetch in flight.""" @@ -184,7 +184,12 @@ class ClientRole: block_ids: Sequence[int], send_ready: bool, ) -> None: - """Register a load request and send the FetchMsg.""" + """Send the FetchMsg closing the current lookup round. + + The scheduler may submit several loads per kv_request_id as its + matched prefix resolves incrementally; each fetch carries the + round it closes so the loads stay independent on the wire. + """ logger.debug( "P2PSession %s: request_blocks job_id=%d kv_request_id=%s " "blocks=%d ready=%s", @@ -195,77 +200,74 @@ class ClientRole: send_ready, ) st = self._get_or_create_request(kv_request_id) - st.load = _InboundLoadState( + round_seq = st.round_seq + st.round_seq += 1 + st.loads[round_seq] = _InboundLoadState( job_id=job_id, submitted_at=time.monotonic(), ) self._active_loads.add(kv_request_id) - st.phase = ClientPhase.FETCH_SENT + st.peer_lookup_open = False self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, FetchMsg.KEYS: list(keys), FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids], + FetchMsg.ROUND_SEQ: round_seq, } ) - # Issuing the fetch ends this request's lookup phase, so drop all - # probe state. Once the peer serves this fetch both sides unpin, so - # the producer may evict the block; a stale cached True would - # otherwise let a re-scheduled lookup() return HIT without - # re-probing, pointing at a block the producer no longer holds. - # Clearing forces a fresh LookupMsg on re-schedule so the producer - # answers from current state. For a symmetric-P2P request (probes - # populated) every fetched block was a confirmed HIT; a PD-only load - # never probes, so probes is empty and the clear is a no-op. - if st.probes: + # Issuing the fetch closes this lookup round, so drop all probe + # state. Once the peer serves the fetch both sides unpin, so a + # stale cached True would let a later lookup() return HIT for a + # block the producer may have evicted; clearing forces a fresh + # probe under the next round. + if st.probed and keys: + assert st.probes, ( + f"symmetric fetch for {kv_request_id} has keys but no probes" + ) assert all(st.probes.get(key) is True for key in keys) st.probes.clear() def finish(self, kv_request_id: str) -> None: - """Finish a request: abort any in-flight load and release lookup state. + """Finish a request: abort in-flight loads and release lookup state. - Called from the session's ``finish_request``. The two branches are - mutually exclusive: ``load`` is set only by ``request_blocks``, which - also advances ``phase`` to ``FETCH_SENT``, and nothing moves it back - to ``PROBING`` — so a fetch in flight never coexists with the - ``PROBING`` phase. - - - Fetch in flight (``load`` set, phase ``FETCH_SENT``): send an - AbortFetchMsg unless the load is already aborting, then drop it. - - Outstanding lookups (``PROBING``): a LookupMsg was flushed but no - FetchMsg has closed the peer's lookup phase. Every FetchMsg the - server receives in p2p mode is its "request finished" signal (it - releases lookup state and fires ``cb.finish_request``); when the - client's lookups all missed no FetchMsg is otherwise sent, so emit - a terminal empty one purely to trigger those semantics. In - ``REGISTERED`` the peer never received a LookupMsg and in - ``FETCH_SENT`` a FetchMsg already closed the phase, so neither owes - a terminal FetchMsg. + Called from the session's ``finish_request``. Sends an + AbortFetchMsg per load not already aborting, and — independently — + the terminal empty FetchMsg when the peer still holds lookup + state no fetch has closed (its "request finished" signal: it + releases lookup state, drains parked supply, and fires + ``cb.finish_request``). A later round's supply can be parked + while an earlier round's load is still in flight, so both can be + owed at once. Then drop all probe/lookup state and prune the entry. """ st = self._requests.get(kv_request_id) if st is None: return - if st.load is not None: - if st.load.aborted_at is None: + if st.loads: + for round_seq, load in st.loads.items(): + if load.aborted_at is not None: + continue self._send( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + AbortFetchMsg.ROUND_SEQ: round_seq, } ) - st.load = None + st.loads.clear() self._active_loads.discard(kv_request_id) - elif st.phase is ClientPhase.PROBING: - st.phase = ClientPhase.FETCH_SENT + if st.peer_lookup_open: + st.peer_lookup_open = False self._send( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.KV_REQUEST_ID: kv_request_id, FetchMsg.KEYS: [], FetchMsg.BLOCK_INDEXES: [], + FetchMsg.ROUND_SEQ: st.round_seq, } ) st.probes.clear() @@ -273,20 +275,21 @@ class ClientRole: self._flush_pending.discard(kv_request_id) self._maybe_prune(kv_request_id) - def on_transfer_done(self, kv_request_id: str, success: bool) -> None: + def on_transfer_done( + self, kv_request_id: str, success: bool, round_seq: int + ) -> None: """Handle a TransferDoneMsg from the peer.""" st = self._requests.get(kv_request_id) - if st is not None and st.load is not None: + load = st.loads.pop(round_seq, None) if st is not None else None + if st is not None and load is not None: self._completed_loads.append( LoadResult( - job_id=st.load.job_id, + job_id=load.job_id, kv_request_id=kv_request_id, success=success, ) ) - st.load = None - self._active_loads.discard(kv_request_id) - self._maybe_prune(kv_request_id) + self._on_load_terminal(kv_request_id, st) else: # No matching in-flight load: either a duplicate # transfer_done from the peer (protocol violation) or a @@ -295,41 +298,44 @@ class ClientRole: # so we can't tell — log so it's findable. logger.warning( "P2PSession %s: transfer_done for unknown kv_request_id=%s " - "(duplicate from peer, or raced with local cancel/timeout)", + "round=%s (duplicate from peer, or raced with local " + "cancel/timeout)", self._peer_id, kv_request_id, + round_seq, ) - def on_abort_ack(self, kv_request_id: str) -> None: + def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None: """Handle an AbortAckMsg from the peer.""" st = self._requests.get(kv_request_id) - if st is not None and st.load is not None: + load = st.loads.pop(round_seq, None) if st is not None else None + if st is not None and load is not None: logger.warning( "P2PSession %s: load request %s (job_id=%d) timed out; " "load job completed with failure. If this recurs, ensure " "PYTHONHASHSEED is set to the same value on all nodes.", self._peer_id, kv_request_id, - st.load.job_id, + load.job_id, ) self._completed_loads.append( LoadResult( - job_id=st.load.job_id, + job_id=load.job_id, kv_request_id=kv_request_id, success=False, ) ) - st.load = None - self._active_loads.discard(kv_request_id) - self._maybe_prune(kv_request_id) + self._on_load_terminal(kv_request_id, st) else: # See on_transfer_done: same ambiguity (duplicate ack # vs. raced with local cancel/timeout that already popped). logger.warning( "P2PSession %s: abort_ack for unknown kv_request_id=%s " - "(duplicate from peer, or raced with local cancel/timeout)", + "round=%s (duplicate from peer, or raced with local " + "cancel/timeout)", self._peer_id, kv_request_id, + round_seq, ) # ------------------------------------------------------------------ @@ -345,17 +351,18 @@ class ClientRole: - Once a LookupRespMsg has resolved the entry: returns the cached bool result on every call without popping it. - A resolved entry is retained until its fetch is issued - (``request_blocks`` pops it) or the request finishes + A resolved entry is retained until a fetch closes the round + (``request_blocks`` clears all probes) or the request finishes (``finish`` clears all entries for the id). A request's block set can be re-probed across steps, so popping on read would make a repeat probe of an already-resolved key look brand-new and re-queue it, emitting a redundant LookupMsg for an answer we already hold. Keeping the entry until fetch makes repeat probes - free; clearing it at fetch forces a fresh probe if the request is - re-scheduled, since the block is unpinned once served. + free; clearing at fetch forces a fresh probe under the next + round, since the block is unpinned once served. """ st = self._get_or_create_request(kv_request_id) + st.probed = True okey = OffloadKey(key) if okey in st.probes: return st.probes[okey] @@ -378,14 +385,11 @@ class ClientRole: ``on_schedule_end()``. A request's block set may be discovered across several scheduler steps, so more than one LookupMsg can go out per kv_request_id — one per step that registered new - keys. register_lookup() de-dups in-flight and already-resolved - (req_id, key) pairs, so each LookupMsg carries only the keys - first probed in that step. The peer's lookup phase for the id is - still closed by exactly one FetchMsg, which the client contract - guarantees is sent after every lookup for the id has resolved - (see request_blocks / finish). Send-gating is handled by - the injected ``_send`` callback (queues until ConnectAckMsg if - needed). + keys, all tagged with the current round. register_lookup() + de-dups in-flight and already-resolved (req_id, key) pairs, so + each LookupMsg carries only the keys first probed in that step. + Send-gating is handled by the injected ``_send`` callback + (queues until ConnectAckMsg if needed). Only requests that registered new keys since the last flush are visited — the ``_flush_pending`` work-list avoids scanning every @@ -395,13 +399,9 @@ class ClientRole: st = self._requests.get(req_id) if st is None or not st.unsent: continue - # Record that the peer now holds lookup state for this id so - # finish knows a terminal empty FetchMsg may be owed. - # Only promote from REGISTERED: once a fetch has gone out - # (FETCH_SENT) a later LookupMsg must not regress the phase, as - # no terminal FetchMsg is owed for an already-fetched request. - if st.phase is ClientPhase.REGISTERED: - st.phase = ClientPhase.PROBING + # The peer now holds lookup state for this id; finish owes a + # terminal empty FetchMsg until a fetch closes it. + st.peer_lookup_open = True logger.debug( "P2P LOOKUP client %s: SEND LookupMsg kv_request_id=%s keys=%d", self._peer_id, @@ -413,6 +413,7 @@ class ClientRole: TYPE_KEY: LookupMsg.TYPE, LookupMsg.KV_REQUEST_ID: req_id, LookupMsg.KEYS: list(st.unsent), + LookupMsg.ROUND_SEQ: st.round_seq, } ) st.unsent = [] @@ -451,52 +452,56 @@ class ClientRole: def collect_results(self) -> list[LoadResult]: """Walk load timeouts and drain completed loads. - Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg - sent and enter the aborting phase. Aborting requests past - ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads. + Loads past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg sent and + enter the aborting phase. Aborting loads past + ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed. Lookups have no timeout: an unanswered probe stays None (RETRY) until finish_request clears it — see ``_ClientRequestState.probes``. """ now = time.monotonic() - to_remove: list[str] = [] + to_remove: list[tuple[str, int]] = [] for req_id in self._active_loads: st = self._requests[req_id] - assert st.load is not None - load = st.load - if load.aborted_at is None: - if now - load.submitted_at >= _LOAD_TIMEOUT_S: - load.aborted_at = now - logger.warning( - "P2PSession %s: %s timed out, sending abort", - self._peer_id, - req_id, - ) - self._send( - { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: req_id, - } - ) - else: - if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: - to_remove.append(req_id) - self._completed_loads.append( - LoadResult( - job_id=load.job_id, - kv_request_id=req_id, - success=False, + assert st.loads + for round_seq, load in st.loads.items(): + if load.aborted_at is None: + if now - load.submitted_at >= _LOAD_TIMEOUT_S: + load.aborted_at = now + logger.warning( + "P2PSession %s: %s round=%s timed out, sending abort", + self._peer_id, + req_id, + round_seq, ) - ) - logger.warning( - "P2PSession %s: abort_ack timed out for kv_request_id=%s", - self._peer_id, - req_id, - ) - for req_id in to_remove: - self._requests[req_id].load = None - self._active_loads.discard(req_id) - self._maybe_prune(req_id) + self._send( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: req_id, + AbortFetchMsg.ROUND_SEQ: round_seq, + } + ) + else: + if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: + to_remove.append((req_id, round_seq)) + self._completed_loads.append( + LoadResult( + job_id=load.job_id, + kv_request_id=req_id, + success=False, + ) + ) + logger.warning( + "P2PSession %s: abort_ack timed out for " + "kv_request_id=%s round=%s", + self._peer_id, + req_id, + round_seq, + ) + for req_id, round_seq in to_remove: + st = self._requests[req_id] + st.loads.pop(round_seq) + self._on_load_terminal(req_id, st) results = self._completed_loads self._completed_loads = [] @@ -510,12 +515,12 @@ class ClientRole: forever on an answer that can never arrive). See ``ClientCloseResult``. """ failed_jobs = [ - st.load.job_id for st in self._requests.values() if st.load is not None + load.job_id for st in self._requests.values() for load in st.loads.values() ] failed_req_ids = [ req_id for req_id, st in self._requests.items() - if st.load is not None or any(hit is None for hit in st.probes.values()) + if st.loads or any(hit is None for hit in st.probes.values()) ] self._requests.clear() self._flush_pending.clear() diff --git a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py index 8988c7e79f3..33ae5bf2f79 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/protocol.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/protocol.py @@ -26,14 +26,12 @@ Block Transfer Flow (happy path) 1. Client sends FetchMsg with a kv_request_id and lists of block keys + remote indexes where it wants the data written. - In p2p mode FetchMsg is also the server-side "request finished" - signal for the id: no further ``cb.create_store_job`` will fire - (parked LookupMsg batches are popped, so pending-key resolution - cannot promote a HIT after this point), all server-side lookup - state for the id is released, and ``cb.finish_request`` fires on - each dropped batch. The client emits exactly one FetchMsg per - lookup-touched request, including an empty one when no blocks - end up being fetched. + A request may run several lookup→fetch rounds; symmetric-P2P + messages carry ROUND_SEQ so each round's supply, demand, and + completion stay isolated. The terminal empty FetchMsg is the + server-side "request finished" signal for the id: parked + LookupMsg batches are popped and ``cb.finish_request`` fires on + each. 2. Server matches requested blocks against locally stored blocks: - Blocks already available are transferred immediately via RDMA. - Blocks not yet available are recorded as "demanded" and @@ -178,33 +176,32 @@ class DisconnectMsg: class FetchMsg: - """Client → Server: request blocks by key and close the lookup phase. + """Client → Server: request blocks for one lookup round. - In p2p mode FetchMsg is also the server-side "request finished" - signal for ``kv_request_id``: on receipt the server (a) fires no - further ``cb.create_store_job`` for this id — parked LookupMsg - batches are popped, so ``_resolve_pending_lookups`` cannot promote - a HIT_PENDING / RETRY key into a fresh pin after this point — and - (b) calls ``cb.finish_request(batch.ctx)`` on each dropped batch - so the TieringManager can release per-batch bookkeeping. In the - all-miss case the client emits an empty FetchMsg (``KEYS`` - and ``BLOCK_INDEXES`` both empty) purely to fire this signal. + A non-empty fetch closes only its round. The terminal empty FetchMsg + (``KEYS`` and ``BLOCK_INDEXES`` both empty) is the "request + finished" signal: the server pops parked LookupMsg batches, calls + ``cb.finish_request`` on each, and drains any leftover supply. Fields: KV_REQUEST_ID: Identifies this block transfer request. KEYS: List of block keys (OffloadKey bytes). May be empty. BLOCK_INDEXES: List of remote block indexes (same length as KEYS). + ROUND_SEQ: Lookup round this fetch closes. PD clients never probe + and stay on their single round 0. """ TYPE = "fetch" KV_REQUEST_ID = "kv_request_id" KEYS = "keys" BLOCK_INDEXES = "block_indexes" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, FetchMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, FetchMsg.ROUND_SEQ) _require_list(msg, FetchMsg.KEYS) _require_list(msg, FetchMsg.BLOCK_INDEXES) keys = msg[FetchMsg.KEYS] @@ -229,17 +226,21 @@ class LookupMsg: Fields: KV_REQUEST_ID: Identifies this lookup transaction. KEYS: List of block keys (OffloadKey bytes) to probe. + ROUND_SEQ: Lookup round these probes belong to; pinned supply is + parked under it for that round's fetch. """ TYPE = "lookup" KV_REQUEST_ID = "kv_request_id" KEYS = "keys" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, LookupMsg.KV_REQUEST_ID, str) _require_list(msg, LookupMsg.KEYS) + _require_non_neg_int(msg, LookupMsg.ROUND_SEQ) class LookupRespMsg: @@ -284,17 +285,22 @@ class TransferDoneMsg: Fields: KV_REQUEST_ID: The request that completed. SUCCESS: Whether the transfer completed successfully. + ROUND_SEQ: The fetch round that completed. Several loads can be + in flight per id (the scheduler submits loads incrementally), + so completions are matched by round. """ TYPE = "transfer_done" KV_REQUEST_ID = "kv_request_id" SUCCESS = "success" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, TransferDoneMsg.KV_REQUEST_ID, str) _require(msg, TransferDoneMsg.SUCCESS, bool) + _require_non_neg_int(msg, TransferDoneMsg.ROUND_SEQ) class AbortFetchMsg: @@ -302,15 +308,18 @@ class AbortFetchMsg: Fields: KV_REQUEST_ID: The request to cancel. + ROUND_SEQ: The fetch round to cancel. """ TYPE = "abort_fetch" KV_REQUEST_ID = "kv_request_id" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, AbortFetchMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, AbortFetchMsg.ROUND_SEQ) class AbortAckMsg: @@ -318,12 +327,15 @@ class AbortAckMsg: Fields: KV_REQUEST_ID: The request that was cancelled. + ROUND_SEQ: The round that was cancelled; echoes AbortFetchMsg. """ TYPE = "abort_ack" KV_REQUEST_ID = "kv_request_id" + ROUND_SEQ = "round_seq" @staticmethod def validate(msg: dict) -> None: """Raise ValueError if any field has an invalid type or value.""" _require(msg, AbortAckMsg.KV_REQUEST_ID, str) + _require_non_neg_int(msg, AbortAckMsg.ROUND_SEQ) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 3709ecd9c53..8558acdf919 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -53,15 +53,6 @@ class StoreResult(NamedTuple): success: bool -class _InflightXfer(NamedTuple): - """Metadata for a single inflight RDMA transfer, keyed by transfer_id.""" - - kv_request_id: str - block_count: int - # The set of store job IDs that contributed blocks to this transfer. - job_ids: set[int] - - class _MatchResult(NamedTuple): """Result of block matching: pairs ready for transfer.""" @@ -73,12 +64,17 @@ class _MatchResult(NamedTuple): @dataclass class _OutboundRequestState: - """Server-role state for a single peer fetch request. + """Server-role state for a single fetch round of a peer request. - The owning ``kv_request_id`` is the ``ServerRole._requests`` dict key - and is not duplicated on the value. + A kv_request_id may run several lookup→fetch rounds; rounds live in + ``_ServerRequestState.outbound`` keyed by the wire ``round_seq``, so + terminals touch only their own round. """ + # Supply came from inbound lookup pins (symmetric): no late + # submit_store can arrive, so unmatched fetch demand fails fast. PD + # rounds park demand for stores instead. + lookup_supplied: bool = False demand_received: bool = False available: dict[OffloadKey, tuple[int, int]] = field( default_factory=dict @@ -88,7 +84,8 @@ class _OutboundRequestState: ) # key → remote_block_idx: blocks peer wants, awaiting supply remaining: int = 0 # blocks that need to be transferred to client finishing: bool = False # Signal finish request ASAP - # Job IDs that submit_store'd blocks for this request and have not + inflight: int = 0 # transfers submitted for this round, not yet polled + # Job IDs that submit_store'd blocks for this round and have not # yet emitted a StoreResult. The terminal-finalize helper drains # this set; poll-done and poll-failed discard entries as their # StoreResults fire. @@ -145,6 +142,21 @@ class _OutboundRequestState: ) +@dataclass +class _InflightXfer: + """Metadata for a single inflight RDMA transfer, keyed by transfer_id.""" + + kv_request_id: str + block_count: int + # The set of store job IDs that contributed blocks to this transfer. + job_ids: set[int] + # Round this transfer serves and its key in ``st.outbound``; + # remaining/finalize apply only while the round is still registered. + # Dummy default for test-seeded entries. + round: _OutboundRequestState = field(default_factory=_OutboundRequestState) + round_key: int = 0 + + @dataclass class _ActiveLookup: """In-flight state for one inbound LookupMsg. @@ -159,6 +171,8 @@ class _ActiveLookup: lookup_id: int kv_request_id: str ctx: ReqContext + # Wire round these probes belong to; pins park under it. + round_seq: int = 0 # Keys from the inbound LookupMsg, preserved in wire order so # the aggregated response goes back in the same order. keys: list[OffloadKey] = field(default_factory=list) @@ -185,6 +199,7 @@ class _PendingLookup(NamedTuple): keys: list[OffloadKey] enqueued_at: float + round_seq: int = 0 @dataclass @@ -198,17 +213,15 @@ class _ServerRequestState: is idle — see ``ServerRole._maybe_prune``. """ - # Outbound serve state (PD + symmetric producer side). None until the - # first add_stored_blocks / on_fetch; reset to None on finalize/abort. - outbound: _OutboundRequestState | None = None + # Fetch rounds keyed by wire round_seq. A round is created by its + # first supply or its fetch and removed at its terminal (finalize / + # failure / abort). + outbound: dict[int, _OutboundRequestState] = field(default_factory=dict) # Raw inbound LookupMsgs not yet processed against the ParentManager. pending_lookups: list[_PendingLookup] = field(default_factory=list) # Per-LookupMsg state parked with HIT_PENDING / RETRY keys, keyed by # the (globally unique) lookup_id and re-polled each serve. lookups: dict[int, _ActiveLookup] = field(default_factory=dict) - # Start time of a pending abort drain (``time.monotonic``); None when - # no abort is in progress. - abort_started_at: float | None = None # Transfer ids in ``ServerRole._inflight`` for this id. Kept in sync # via _inflight_add / _inflight_pop so a non-empty set is an exact # "has any inflight transfer" predicate and the abort drain can @@ -256,9 +269,9 @@ class ServerRole: # ``parent.on_request_finished`` in ``serve_external_requests``. self._finished_lookup_ctxs: list[ReqContext] = [] self._lookup_id_counter: int = 0 - # kv_request_ids with a parked abort awaiting drain — work-list so - # drain_pending_aborts doesn't scan every request each poll tick. - self._parked_aborts: set[str] = set() + # Parked aborts awaiting drain, keyed by (kv_request_id, round) + # with the abort start time. + self._pending_aborts: dict[tuple[str, int], float] = {} # ------------------------------------------------------------------ # State helpers @@ -277,11 +290,11 @@ class ServerRole: st = self._requests.get(kv_request_id) if ( st is not None - and st.outbound is None + and not st.outbound and not st.inflight_tids and not st.lookups and not st.pending_lookups - and st.abort_started_at is None + and not any(kv == kv_request_id for kv, _ in self._pending_aborts) ): del self._requests[kv_request_id] @@ -295,87 +308,100 @@ class ServerRole: keys: Sequence[OffloadKey], block_ids: Sequence[int], job_id: JobId, + round_seq: int = 0, + *, + from_lookup: bool = False, ) -> None: - """New blocks stored locally — match against pending fetch demand.""" + """New blocks stored locally — match within their fetch round. + + Lookup pins carry the round they were probed under; PD + submit_store batches share PD's single round 0. + """ self._store_jobs[job_id] = time.monotonic() st = self._get_or_create_request(kv_request_id) - if st.outbound is None: - st.outbound = _OutboundRequestState() - result = st.outbound.add_stored_blocks(keys, block_ids, job_id) - if result.local_idxs and st.outbound.demand_received: - self._submit_transfer(kv_request_id, result) + rnd = st.outbound.get(round_seq) + if rnd is None: + rnd = st.outbound[round_seq] = _OutboundRequestState() + if from_lookup: + rnd.lookup_supplied = True + result = rnd.add_stored_blocks(keys, block_ids, job_id) + if result.local_idxs and rnd.demand_received: + self._submit_transfer(kv_request_id, result, rnd, round_seq) def on_fetch( self, kv_request_id: str, keys: Sequence[OffloadKey], block_indexes: Sequence[int], + round_seq: int = 0, ) -> None: """Handle a FetchMsg from the peer. - In p2p mode FetchMsg is the server-side "request finished" - signal for ``kv_request_id``. Three consequences flow from that: - - - No further ``parent.create_store_job`` will fire for this id: - the request's parked ``lookups`` are popped here (via - ``_finish_inbound_lookups``) before the next - ``serve_external_requests`` runs ``_resolve_pending_lookups``, - so any HIT_PENDING / RETRY key that would otherwise later - promote to HIT and pin a slot is dropped instead. Any raw - not-yet-processed LookupMsg for this id is dropped too. - - All server-side lookup state for the id is cleaned up (the - ``lookups`` entries themselves). - - The synthetic ctx is queued for ``parent.on_request_finished`` - (fired by the next ``serve_external_requests``) so the - TieringManager can release per-lookup bookkeeping. - - The client contract guarantees exactly one FetchMsg per - lookup-touched request — including an empty one when no - blocks end up being fetched. - - Raises ``ValueError`` on a duplicate fetch for the same - ``kv_request_id``; the coordinator's dispatch loop turns that - into a protocol-error disconnect. + A non-empty fetch binds and closes its round, leaving lookup + state alone (the next round's LookupMsg may already be in + flight). The terminal empty fetch closes the id: parked lookups + are popped and every remaining round drained. A second fetch for + a round already holding demand raises ValueError + (protocol-error disconnect). """ logger.debug( - "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d", + "P2PSession %s: fetch RECEIVED kv_request_id=%s round=%s blocks=%d", self._peer_id, kv_request_id, + round_seq, len(keys), ) st = self._requests.get(kv_request_id) - existing = st.outbound if st is not None else None + existing = st.outbound.get(round_seq) if st is not None else None if existing is not None and existing.demand_received: - # A second fetch for the same kv_request_id would overwrite - # `remaining` and leak inflight bookkeeping. Treat as a - # protocol violation. - raise ValueError(f"duplicate fetch for kv_request_id={kv_request_id}") + raise ValueError( + f"duplicate fetch for kv_request_id={kv_request_id} round={round_seq}" + ) st = self._get_or_create_request(kv_request_id) - if st.outbound is None: - st.outbound = _OutboundRequestState() - req = st.outbound + req = st.outbound.get(round_seq) + if req is None: + req = st.outbound[round_seq] = _OutboundRequestState() result = req.add_fetch_demand(keys, block_indexes) + if not keys: + # Terminal empty fetch: close the lookup phase and drain + # every round with no TransferDoneMsg (nothing waits on it). + self._finish_inbound_lookups(kv_request_id) + for key in list(st.outbound): + self._finalize_outbound(kv_request_id, key, send_done=False) + return + if req.lookup_supplied and req.demanded: + # A symmetric round's supply always precedes its fetch, so + # unmatched demand is unservable — fail now, not at the load + # timeout. PD rounds keep parking demand for stores that + # arrive later. + logger.warning( + "P2PSession %s: fetch kv_request_id=%s round=%s demanded %d " + "blocks but %d have no pinned supply; failing fetch " + "immediately", + self._peer_id, + kv_request_id, + round_seq, + len(keys), + len(req.demanded), + ) + self._finalize_outbound(kv_request_id, round_seq, success=False) + return if result.local_idxs: - self._submit_transfer(kv_request_id, result) - # Close the peer's request as far as the server's lookup phase - # is concerned: pop parked lookups so no further - # ``parent.create_store_job`` fires for this id, and queue their - # ctxs for ``parent.on_request_finished``. Done before the - # finalize path below so bookkeeping releases in-order. - self._finish_inbound_lookups(kv_request_id) + self._submit_transfer(kv_request_id, result, req, round_seq) # Prefiller-first mode: finish_request may have run before # fetch arrived. If so, finalize once we know what was # demanded — fully satisfied → success, else early-fail. - if req.finishing and not self._has_inflight_for(kv_request_id): - self._finalize_outbound(kv_request_id) + if req.finishing and req.inflight == 0: + self._finalize_outbound(kv_request_id, round_seq) - def on_abort_fetch(self, kv_request_id: str) -> None: - """Handle an AbortFetchMsg from the peer.""" + def on_abort_fetch(self, kv_request_id: str, round_seq: int = 0) -> None: + """Handle an AbortFetchMsg from the peer, cancelling one round.""" # Abort for an unknown id may be a benign race/duplicate or a # real protocol violation; we don't track completed ids, so warn. st = self._requests.get(kv_request_id) - has_outbound = st is not None and st.outbound is not None - if not has_outbound and not self._has_inflight_for(kv_request_id): + if (st is None or not st.outbound) and not self._has_inflight_for( + kv_request_id + ): logger.warning( "P2PSession %s: abort_fetch for unknown kv_request_id=%s " "(no outbound or inflight state); benign race or stale", @@ -385,16 +411,15 @@ class ServerRole: # Idempotent: receiving AbortFetchMsg again before we've sent the # ack just triggers another drain attempt without resetting the # deadline. - st = self._get_or_create_request(kv_request_id) - if st.abort_started_at is None: - st.abort_started_at = time.monotonic() - self._parked_aborts.add(kv_request_id) - self._drain_abort(kv_request_id) + self._get_or_create_request(kv_request_id) + self._pending_aborts.setdefault((kv_request_id, round_seq), time.monotonic()) + self._drain_abort(kv_request_id, round_seq) def on_lookup( self, kv_request_id: str, keys: Sequence[OffloadKey], + round_seq: int = 0, ) -> None: """Enqueue a LookupMsg from a symmetric-P2P consumer. @@ -406,13 +431,18 @@ class ServerRole: parent calls are valid. """ logger.debug( - "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s keys=%d", + "P2P LOOKUP server %s: RECV LookupMsg kv_request_id=%s round=%s keys=%d", self._peer_id, kv_request_id, + round_seq, len(keys), ) self._get_or_create_request(kv_request_id).pending_lookups.append( - _PendingLookup(keys=list(keys), enqueued_at=time.monotonic()) + _PendingLookup( + keys=list(keys), + enqueued_at=time.monotonic(), + round_seq=round_seq, + ) ) self._serve_pending.add(kv_request_id) @@ -434,7 +464,7 @@ class ServerRole: st.pending_lookups = [] for pl in pending: self._process_inbound_lookup( - kv_request_id, pl.keys, pl.enqueued_at, parent + kv_request_id, pl.keys, pl.enqueued_at, pl.round_seq, parent ) self._resolve_pending_lookups(kv_request_id, parent) st = self._requests.get(kv_request_id) @@ -481,9 +511,7 @@ class ServerRole: else: lookup.pending.add(h) if new_hits: - self._pin_and_register_hits( - lookup.kv_request_id, new_hits, lookup.ctx, parent - ) + self._pin_and_register_hits(lookup, new_hits, parent) return new_hits def _process_inbound_lookup( @@ -491,6 +519,7 @@ class ServerRole: kv_request_id: str, keys: list[OffloadKey], enqueued_at: float, + round_seq: int, parent: ParentManager, ) -> None: """Resolve one enqueued LookupMsg against ``parent``. @@ -515,6 +544,7 @@ class ServerRole: lookup_id=lookup_id, kv_request_id=kv_request_id, ctx=ctx, + round_seq=round_seq, keys=list(keys), deadline=enqueued_at + _LOOKUP_PENDING_TIMEOUT_S, ) @@ -547,25 +577,26 @@ class ServerRole: def _pin_and_register_hits( self, - kv_request_id: str, + lookup: _ActiveLookup, keys: list[OffloadKey], - ctx: ReqContext, parent: ParentManager, ) -> None: - """Pin primary slots for HIT keys and feed them into the - existing ``add_stored_blocks`` matching path. + """Pin primary slots for HIT keys and park them as the lookup's + round supply via ``add_stored_blocks``. Caller has already confirmed every key is HIT (single-threaded scheduler ⇒ no eviction race), so the JobMetadata returned by ``parent.create_store_job`` carries parallel ``keys``/``block_ids`` of length ``len(keys)``. """ - meta = parent.create_store_job(keys, ctx) + meta = parent.create_store_job(keys, lookup.ctx) self.add_stored_blocks( - kv_request_id, + lookup.kv_request_id, list(meta.keys), list(meta.block_ids), meta.job_id, + round_seq=lookup.round_seq, + from_lookup=True, ) def _resolve_pending_lookups( @@ -650,9 +681,8 @@ class ServerRole: Called on the two events that mean "no more lookup traffic for ``kv_request_id`` is expected on this session": the terminal - FetchMsg from the peer (client contract: exactly one FetchMsg - per lookup-touched request, even if empty), and a local - ``finish``. Whichever fires second is a no-op. + empty FetchMsg from the peer and a local ``finish``. Whichever + fires second is a no-op. """ st = self._requests.get(kv_request_id) if st is None: @@ -688,20 +718,15 @@ class ServerRole: self._finish_inbound_lookups(kv_request_id) st = self._requests.get(kv_request_id) - req = st.outbound if st is not None else None - if req is None: + if st is None: return - req.finishing = True - if not req.demand_received: - return - if self._has_inflight_for(kv_request_id): - return - # Remaining > 0 here: if it had hit 0, the poll-done success - # branch would have already cleared outbound and we'd have - # returned at `req is None` above. Helper derives success from - # remaining and emits StoreResult(success=False) for any - # leftover pending jobs. - self._finalize_outbound(kv_request_id) + for key, req in list(st.outbound.items()): + req.finishing = True + if not req.demand_received or req.inflight: + # No demand yet (prefiller-first): on_fetch finalizes via + # `finishing`. Inflight: the last completion finalizes. + continue + self._finalize_outbound(kv_request_id, key) def collect_results(self) -> list[StoreResult]: """Drain timeouts, deferred results, and transport completions. @@ -741,20 +766,24 @@ class ServerRole: ) continue results.extend(self._settle_xfer_jobs(xfer, success=True)) + rnd = xfer.round st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None - if req is not None and req.demand_received: - req.remaining -= xfer.block_count - assert req.remaining >= 0, ( + if st is not None and st.outbound.get(xfer.round_key) is rnd: + rnd.remaining -= xfer.block_count + assert rnd.remaining >= 0, ( f"remaining went negative for kv_request_id={xfer.kv_request_id}" ) - if req.remaining == 0: - self._finalize_outbound(xfer.kv_request_id, success=True) - elif req.finishing and not self._has_inflight_for(xfer.kv_request_id): - self._finalize_outbound(xfer.kv_request_id, success=False) + if rnd.remaining == 0: + self._finalize_outbound( + xfer.kv_request_id, xfer.round_key, success=True + ) + elif rnd.finishing and rnd.inflight == 0: + self._finalize_outbound( + xfer.kv_request_id, xfer.round_key, success=False + ) self._maybe_prune(xfer.kv_request_id) - failed_kv_request_ids: set[str] | None = None + failed_rounds: list[tuple[str, _OutboundRequestState]] | None = None for tid in poll_result.failed: xfer = self._inflight_pop(tid) if xfer is None: @@ -767,35 +796,36 @@ class ServerRole: tid, ) continue - if failed_kv_request_ids is None: - failed_kv_request_ids = set() - failed_kv_request_ids.add(xfer.kv_request_id) results.extend(self._settle_xfer_jobs(xfer, success=False)) + rnd = xfer.round st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None - if st is not None: - st.outbound = None - if req is not None and req.demand_received: + if st is not None and st.outbound.get(xfer.round_key) is rnd: + del st.outbound[xfer.round_key] + if failed_rounds is None: + failed_rounds = [] + failed_rounds.append((xfer.kv_request_id, rnd)) self._send( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id, TransferDoneMsg.SUCCESS: False, + TransferDoneMsg.ROUND_SEQ: xfer.round_key, } ) self._maybe_prune(xfer.kv_request_id) - # Cancel other inflight for the same failed kv_request_ids - if failed_kv_request_ids: - ids_to_cancel = [ - tid - for tid, xfer in self._inflight.items() - if xfer.kv_request_id in failed_kv_request_ids - ] - for tid in ids_to_cancel: - self._inflight_pop(tid) - self._transport.cancel(ids_to_cancel) - for kv_request_id in failed_kv_request_ids: + # Cancel each failed round's other inflight and fail its + # remaining store jobs — nothing else will settle them. + if failed_rounds: + for kv_request_id, rnd in failed_rounds: + ids_to_cancel = [ + tid for tid, x in self._inflight.items() if x.round is rnd + ] + for tid in ids_to_cancel: + self._inflight_pop(tid) + if ids_to_cancel: + self._transport.cancel(ids_to_cancel) + results.extend(self._fail_round_jobs(rnd)) self._maybe_prune(kv_request_id) return results @@ -811,8 +841,8 @@ class ServerRole: def drain_pending_aborts(self) -> None: """Re-attempt every parked abort once per poll tick.""" - for kv_request_id in list(self._parked_aborts): - self._drain_abort(kv_request_id) + for kv_request_id, round_seq in list(self._pending_aborts): + self._drain_abort(kv_request_id, round_seq) def close(self) -> tuple[list[int], list[ReqContext]]: """Tear down. Cancels inflight. @@ -839,7 +869,7 @@ class ServerRole: failed_serves.extend(self._finished_lookup_ctxs) self._requests.clear() self._serve_pending.clear() - self._parked_aborts.clear() + self._pending_aborts.clear() self._finished_lookup_ctxs.clear() return failed_stores, failed_serves @@ -870,6 +900,8 @@ class ServerRole: xfer = self._inflight.pop(tid, None) if xfer is None: return None + xfer.round.inflight -= 1 + assert xfer.round.inflight >= 0 st = self._requests.get(xfer.kv_request_id) if st is not None: st.inflight_tids.discard(tid) @@ -880,80 +912,108 @@ class ServerRole: ) -> list[StoreResult]: """Emit StoreResults for a completed transfer's store jobs. - Pops each attached job from ``_store_jobs`` and clears it from the - request's pending set. A job already popped (via timeout, cancel, - etc.) is skipped so we never double-emit a contradictory result. + Pops each attached job from ``_store_jobs`` and clears it from + its round's pending set. A job already popped (via timeout, + cancel, etc.) is skipped so we never double-emit a contradictory + result. """ results: list[StoreResult] = [] - st = self._requests.get(xfer.kv_request_id) - req = st.outbound if st is not None else None for job_id in xfer.job_ids: if self._store_jobs.pop(job_id, None) is None: continue results.append(StoreResult(job_id=job_id, success=success)) - if req is not None: - req.pending_job_ids.discard(job_id) + xfer.round.pending_job_ids.discard(job_id) return results # ------------------------------------------------------------------ # Internal — finalize / abort drain # ------------------------------------------------------------------ + def _fail_round_jobs(self, rnd: _OutboundRequestState) -> list[StoreResult]: + """Fail a terminated round's still-pending store jobs (idempotent).""" + results: list[StoreResult] = [] + for job_id in rnd.pending_job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + results.append(StoreResult(job_id=job_id, success=False)) + rnd.pending_job_ids.clear() + return results + def _finalize_outbound( self, kv_request_id: str, + round_key: int, success: bool | None = None, + send_done: bool = True, ) -> None: - """Pop the outbound state and emit terminal results. - - Called when no further work will happen for this kv_request_id - on the server side: either request_finish has fired and there - are no inflight transfers, or the last inflight just completed - while finishing. + """Pop one round and emit its terminal results. If ``success`` is None, derive it from ``req.remaining == 0``. - The same flag is used for both the peer's TransferDoneMsg and - the StoreResult(s) emitted for any leftover pending job_ids. + ``send_done=False`` skips the TransferDoneMsg (terminal empty + fetch). Other rounds of the id are untouched. """ st = self._requests[kv_request_id] - assert st.outbound is not None - req = st.outbound - st.outbound = None + req = st.outbound.pop(round_key) if success is None: - success = req.remaining == 0 - for job_id in req.pending_job_ids: - self._store_jobs.pop(job_id, None) - self._pending_store_results.append( - StoreResult(job_id=job_id, success=success) - ) - self._send( - { - TYPE_KEY: TransferDoneMsg.TYPE, - TransferDoneMsg.KV_REQUEST_ID: kv_request_id, - TransferDoneMsg.SUCCESS: success, - } + success = req.demand_received and req.remaining == 0 + settled = self._fail_round_jobs(req) if not success else None + if settled is not None: + self._pending_store_results.extend(settled) + else: + for job_id in req.pending_job_ids: + if self._store_jobs.pop(job_id, None) is None: + continue + self._pending_store_results.append( + StoreResult(job_id=job_id, success=True) + ) + req.pending_job_ids.clear() + logger.debug( + "P2PSession %s: finalize kv_request_id=%s round=%s success=%s " + "remaining=%d leftover_available=%d send_done=%s", + self._peer_id, + kv_request_id, + round_key, + success, + req.remaining, + len(req.available), + send_done, ) + if send_done and req.demand_received: + self._send( + { + TYPE_KEY: TransferDoneMsg.TYPE, + TransferDoneMsg.KV_REQUEST_ID: kv_request_id, + TransferDoneMsg.SUCCESS: success, + TransferDoneMsg.ROUND_SEQ: round_key, + } + ) self._maybe_prune(kv_request_id) - def _drain_abort(self, kv_request_id: str) -> None: + def _drain_abort(self, kv_request_id: str, round_seq: int) -> None: """One drain attempt for a pending abort. - Stops accepting more blocks for ``kv_request_id``, then asks the - transport to cancel any matching inflight transfers in - ``mode="wait"``. Sends ``AbortAckMsg`` once nothing remains - inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to - ``mode="immediate"`` and acks anyway. + Detaches the aborted round, then asks the transport to cancel its + inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` once + nothing remains inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` + falls back to ``mode="immediate"`` and acks anyway. """ st = self._requests[kv_request_id] - st.outbound = None - ids = list(st.inflight_tids) + rnd = st.outbound.pop(round_seq, None) + if rnd is not None: + # Its transfers are being cancelled; fail its jobs now + # instead of leaking them to the store timeout. + self._pending_store_results.extend(self._fail_round_jobs(rnd)) + ids = [ + tid + for tid, x in self._inflight.items() + if x.kv_request_id == kv_request_id and x.round_key == round_seq + ] if not ids: - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) return - assert st.abort_started_at is not None - expired = time.monotonic() - st.abort_started_at >= _CANCEL_DRAIN_TIMEOUT_S - if expired: + started_at = self._pending_aborts[(kv_request_id, round_seq)] + if time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S: for tid in ids: self._inflight_pop(tid) self._transport.cancel(ids, mode="immediate") @@ -964,7 +1024,7 @@ class ServerRole: kv_request_id, len(ids), ) - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) return still = self._transport.cancel(ids, mode="wait") @@ -977,16 +1037,15 @@ class ServerRole: if tid not in still_set: self._inflight_pop(tid) if not still: - self._finalize_abort(kv_request_id) + self._finalize_abort(kv_request_id, round_seq) - def _finalize_abort(self, kv_request_id: str) -> None: - st = self._requests[kv_request_id] - st.abort_started_at = None - self._parked_aborts.discard(kv_request_id) + def _finalize_abort(self, kv_request_id: str, round_seq: int) -> None: + self._pending_aborts.pop((kv_request_id, round_seq), None) self._send( { TYPE_KEY: AbortAckMsg.TYPE, AbortAckMsg.KV_REQUEST_ID: kv_request_id, + AbortAckMsg.ROUND_SEQ: round_seq, } ) self._maybe_prune(kv_request_id) @@ -995,7 +1054,13 @@ class ServerRole: # Internal — transfers and store-job timeouts # ------------------------------------------------------------------ - def _submit_transfer(self, kv_request_id: str, result: _MatchResult) -> None: + def _submit_transfer( + self, + kv_request_id: str, + result: _MatchResult, + rnd: _OutboundRequestState, + round_key: int, + ) -> None: logger.debug( "P2PSession %s: NIXL write_blocks CALL kv_request_id=%s " "local_idxs=%d remote_idxs=%d", @@ -1016,12 +1081,15 @@ class ServerRole: transfer_id, len(result.local_idxs), ) + rnd.inflight += 1 self._inflight_add( transfer_id, _InflightXfer( kv_request_id=kv_request_id, block_count=len(result.local_idxs), job_ids=result.job_ids, + round=rnd, + round_key=round_key, ), ) else: @@ -1031,22 +1099,24 @@ class ServerRole: kv_request_id, len(result.local_idxs), ) - # The matched blocks were popped from req.demanded / - # req.available, but no inflight will satisfy them, so - # remaining will never reach 0 on its own. Mark the - # request as finishing so the existing terminal paths - # clean up: if other inflight is in flight, the last one - # to drain will fire _finalize_outbound(success=False) - # via the elif branch in collect_results. If - # nothing else is in flight, finalize now so the peer - # and the local store jobs don't wait for finish_request - # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S. + # The matched blocks were popped from rnd.demanded / + # rnd.available, but no inflight will satisfy them, so + # remaining will never reach 0 on its own. Mark the round + # as finishing so the existing terminal paths clean up: if + # other transfers of this round are in flight, the last one + # to drain will fire _finalize_outbound(success=False) via + # the elif branch in collect_results. If nothing else is in + # flight, finalize now so the peer and the local store jobs + # don't wait for finish_request or for _STORE_TIMEOUT_S / + # _LOAD_TIMEOUT_S. + rnd.finishing = True st = self._requests.get(kv_request_id) - req = st.outbound if st is not None else None - if req is not None: - req.finishing = True - if not self._has_inflight_for(kv_request_id): - self._finalize_outbound(kv_request_id, success=False) + if ( + st is not None + and st.outbound.get(round_key) is rnd + and rnd.inflight == 0 + ): + self._finalize_outbound(kv_request_id, round_key, success=False) def _timeout_pending_store_jobs(self) -> list[StoreResult]: if not self._store_jobs: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index 7d19913b0a2..4bd780159a4 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -374,26 +374,34 @@ class P2PSession: for bh in msg[FetchMsg.KEYS] ] block_indexes = msg[FetchMsg.BLOCK_INDEXES] + round_seq = msg[FetchMsg.ROUND_SEQ] # Run the server-role state machine inline as today — # add_fetch_demand records demand against any blocks we've # already seen in `available`. Report the kv_request_id so # the manager (after poll() returns) can replay any parked # submit_store batches; their add_stored_blocks calls hit # the demand recorded here and submit transfers immediately. - self._server.on_fetch(kv_request_id, keys, block_indexes) + self._server.on_fetch(kv_request_id, keys, block_indexes, round_seq) self._new_fetch_ids.append(kv_request_id) elif msg_type == AbortFetchMsg.TYPE: AbortFetchMsg.validate(msg) - self._server.on_abort_fetch(msg[AbortFetchMsg.KV_REQUEST_ID]) + self._server.on_abort_fetch( + msg[AbortFetchMsg.KV_REQUEST_ID], + msg[AbortFetchMsg.ROUND_SEQ], + ) elif msg_type == TransferDoneMsg.TYPE: TransferDoneMsg.validate(msg) self._client.on_transfer_done( msg[TransferDoneMsg.KV_REQUEST_ID], msg[TransferDoneMsg.SUCCESS], + msg[TransferDoneMsg.ROUND_SEQ], ) elif msg_type == AbortAckMsg.TYPE: AbortAckMsg.validate(msg) - self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID]) + self._client.on_abort_ack( + msg[AbortAckMsg.KV_REQUEST_ID], + msg[AbortAckMsg.ROUND_SEQ], + ) elif msg_type == LookupMsg.TYPE: LookupMsg.validate(msg) kv_request_id = msg[LookupMsg.KV_REQUEST_ID] @@ -401,7 +409,7 @@ class P2PSession: OffloadKey(bh if isinstance(bh, bytes) else bytes(bh)) for bh in msg[LookupMsg.KEYS] ] - self._server.on_lookup(kv_request_id, keys) + self._server.on_lookup(kv_request_id, keys, msg[LookupMsg.ROUND_SEQ]) elif msg_type == LookupRespMsg.TYPE: LookupRespMsg.validate(msg) kv_request_id = msg[LookupRespMsg.KV_REQUEST_ID] From 4fb483ca86566d163a886416cdd822cf716b1df5 Mon Sep 17 00:00:00 2001 From: IBRAHIM IBRAHIM <66755652+Ibrahim2595@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:47:56 -0400 Subject: [PATCH 32/67] [Docs] Expand llm-d integration page (#45432) Signed-off-by: ibrahimibrahim Co-authored-by: ibrahimibrahim Co-authored-by: Claude Opus 4.8 (1M context) --- docs/deployment/integrations/llm-d.md | 36 +++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/docs/deployment/integrations/llm-d.md b/docs/deployment/integrations/llm-d.md index 6060b98f642..7d261eb910b 100644 --- a/docs/deployment/integrations/llm-d.md +++ b/docs/deployment/integrations/llm-d.md @@ -1,5 +1,37 @@ # llm-d -vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers. +[llm-d](https://llm-d.ai/) is a Kubernetes-native distributed inference framework for serving large language models at scale, with vLLM as its primary inference engine. llm-d coordinates a fleet of vLLM instances across a cluster so that performance holds up under real production traffic, achieving the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators. -You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). +It is a [CNCF Sandbox project](https://www.cncf.io/blog/2026/03/24/welcome-llm-d-to-the-cncf-evolving-kubernetes-into-sota-ai-infrastructure/) founded by Red Hat, Google Cloud, IBM Research, CoreWeave, and NVIDIA. + +## What llm-d adds to vLLM + +A single vLLM server is fast, but at scale the picture changes: across many replicas, cache locality breaks under round-robin load balancing, long prompts inflate time-to-first-token, and accelerators sit underused. llm-d adds the cluster-level layer that vLLM does not aim to provide on its own: + +- **[Prefix-aware routing](https://llm-d.ai/docs/guides/precise-prefix-cache-aware).** Instead of round-robin, llm-d reads vLLM's KV-cache events and routes each request to the replica that already holds its prefix, reusing cache instead of recomputing it. +- **[Distributed KV-cache management](https://llm-d.ai/docs/guides#advanced-kv-cache-management).** A global index tracks which token blocks live on which replica, and [tiered offloading](https://llm-d.ai/docs/guides/tiered-prefix-cache) spills cache to CPU memory or local SSD, extending the working set beyond accelerator HBM. +- **[Prefill/decode disaggregation](https://llm-d.ai/docs/guides/pd-disaggregation).** Prompt processing and token generation run on separate vLLM workers, with KV-cache moved over the vLLM [NIXL connector](https://docs.vllm.ai/en/latest/features/nixl_connector_usage/), lowering TTFT and steadying per-token latency on long prompts. +- **[Wide expert-parallelism](https://llm-d.ai/docs/guides/wide-expert-parallelism).** Serve large Mixture-of-Experts models such as DeepSeek-R1 and GPT-OSS across nodes with combined data and expert parallelism, for more KV-cache capacity and throughput. +- **SLO-aware [autoscaling](https://llm-d.ai/docs/guides/workload-autoscaling) and [flow control](https://llm-d.ai/docs/guides/flow-control).** Scale vLLM pools on real inference signals (queue depth, true demand) rather than raw GPU utilization, with multi-tenant fairness and priority dispatch. + +These are composable. Most teams start by adding prefix-aware routing over an existing vLLM pool, then layer in the rest as specific bottlenecks appear. + +## Performance + +Representative benchmarked results across accelerators: + +- **3x higher output throughput** and **2x faster TTFT** from prefix-aware routing vs round-robin (Llama 3.1 70B, AMD MI300X) +- **Up to 70% higher tokens/sec** from prefill/decode disaggregation (GPT-OSS, NVIDIA B200) +- **13.9x throughput** from hierarchical KV offloading at high concurrency vs GPU-only (NVIDIA H100) + +See the [full list](https://github.com/llm-d/llm-d#performance-highlights) and reproducible benchmarks on [Prism](https://prism.llm-d.ai/). + +## Get started + +1. Deploy the [Optimized Baseline](https://llm-d.ai/docs/guides/optimized-baseline) with the [Quickstart](https://llm-d.ai/docs/getting-started/quickstart). It stands up an intelligent router over a vLLM pool on Kubernetes in a tested configuration. +2. Browse the [well-lit path guides](https://llm-d.ai/docs/guides), each a tested recipe for one of the capabilities above, and add the optimization that fits your workload. +3. Read the [Introduction](https://llm-d.ai/docs/getting-started) and [Architecture overview](https://llm-d.ai/docs/architecture) to see how the pieces wrap your vLLM deployment. + +You can also deploy vLLM with llm-d via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview). + +Questions and contributions are welcome on [GitHub](https://github.com/llm-d/llm-d) and [Slack](https://llm-d.ai/slack). From 6453fc0b8cb50668dceaf5839033b0a2a8fa7f7f Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 09:09:28 -0700 Subject: [PATCH 33/67] [Bugfix] Don't reuse engine core payload buffer while zmq is sending it (#50053) Signed-off-by: Nick Hill Co-authored-by: Andreas Karatzas --- tests/v1/test_serial_utils.py | 136 ++++++++++++++++++++++++++++++++++ vllm/envs.py | 2 +- vllm/v1/engine/core.py | 41 +++++++--- vllm/v1/engine/core_client.py | 58 +++------------ 4 files changed, 179 insertions(+), 58 deletions(-) diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 4ed8724e60f..9f7761c22b5 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc(): assert torch.allclose(decoded.prompt_embeds, original_tensor), ( f"Value mismatch for sender {sender_idx} msg {msg_idx}" ) + + +def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): + """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends + it: many requests, each with per-token tensors small enough that pyzmq + copies their frames, while the accumulated payload frame is large enough + that pyzmq sends it zero-copy.""" + from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs + from vllm.v1.outputs import LogprobsTensors + + outputs = [] + for req in range(num_reqs): + num_tokens = num_prompt_tokens + req % 4 + outputs.append( + EngineCoreOutput( + request_id=f"req-{req:08d}", + new_token_ids=[req], + new_prompt_logprobs_tensors=LogprobsTensors( + logprob_token_ids=torch.arange( + num_tokens * 2, dtype=torch.int64 + ).view(num_tokens, 2), + logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32), + selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32), + ), + ) + ) + return EngineCoreOutputs(outputs=outputs) + + +def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages(): + """The engine core recycles the msgpack payload buffer across messages + (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished + sending that buffer, otherwise a newer payload is delivered alongside the + older message's zero-copy tensor frames. + + `Socket.send_multipart(track=True)` cannot be used to detect this: it + returns a tracker for the last frame only, and pyzmq copies frames below + `zmq.COPY_THRESHOLD` and reports them as already-sent. + """ + import zmq + + from vllm.v1.engine import EngineCoreOutputs + from vllm.v1.engine.core import EngineCoreProc + + num_msgs = 100 + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(EngineCoreOutputs) + # Enough requests that the payload frame is zero-copied rather than copied + # by pyzmq, which is what makes early reuse observable. + messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)] + assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD + + reuse_buffers: list[bytearray] = [] + pending: list[tuple[zmq.MessageTracker, bytearray]] = [] + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-payload-reuse") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-payload-reuse") + + for outputs in messages: + while pending and pending[0][0].done: + reuse_buffers.append(pending.pop(0)[1]) + buffer = reuse_buffers.pop() if reuse_buffers else bytearray() + buffers = encoder.encode_into(outputs, buffer) + tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers) + if tracker.done: + reuse_buffers.append(buffer) + else: + pending.append((tracker, buffer)) + + for i, sent in enumerate(messages): + received = decoder.decode(pull.recv_multipart(copy=False)) + assert len(received.outputs) == len(sent.outputs), f"message {i}" + for expected, actual in zip(sent.outputs, received.outputs): + sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids + got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids + assert actual.request_id == expected.request_id, f"message {i}" + assert torch.equal(got_ids, sent_ids), ( + f"message {i} request {actual.request_id}: corrupted " + f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}" + ) + push.close(linger=0) + pull.close(linger=0) + + +def test_zero_copy_frames_survive_without_caller_side_references(): + """Callers don't need to retain the encoded object until zmq has sent it: + for a zero-copy frame, zmq holds its own reference to the backing buffer. + + The engine core clients rely on this when sending requests that carry + tensors (e.g. prompt embeds) without tracking the messages. + + What makes that safe is that `tensor_data()` hands zmq a memoryview which + transitively references the source tensor, so refcounting - not timing - + keeps the memory from being freed and reused underneath zmq. + """ + import gc + + import zmq + + from vllm.v1.utils import tensor_data + + num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD + expected = torch.arange(num_elems, dtype=torch.int64) + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(RequestWithTensor) + + # The buffer handed to zmq must keep the tensor's storage alive by itself. + holder = tensor_data(expected).obj + while getattr(holder, "base", None) is not None: + holder = holder.base + assert isinstance(holder, torch.Tensor) + assert holder.data_ptr() == expected.data_ptr() + + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-zero-copy-lifetime") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-zero-copy-lifetime") + + request = RequestWithTensor(prompt_embeds=expected.clone(), data="req") + buffers = encoder.encode(request) + assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD + push.send_multipart(buffers, copy=False) + + # Drop every reference the sender holds, then churn the allocator. + del request, buffers + gc.collect() + torch.arange(num_elems * 4, dtype=torch.int64) + + decoded = decoder.decode(pull.recv_multipart(copy=False)) + assert decoded.prompt_embeds is not None + assert torch.equal(decoded.prompt_embeds, expected) + push.close(linger=0) + pull.close(linger=0) diff --git a/vllm/envs.py b/vllm/envs.py index fb54619c748..f7e01c33275 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1547,7 +1547,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # tensors above will instead be sent via a separate message. # While the sending side still actually copies the tensor # in all cases, on the receiving side, tensors above this - # limit will actually be zero-copy decoded. + # limit will actually be zero-copy decoded. The unit is bytes. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int( os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256") ), diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ecac92f5fe0..66135273002 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -7,7 +7,7 @@ import signal import threading import time from collections import defaultdict, deque -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Sequence from concurrent.futures import Future from contextlib import ExitStack, contextmanager from enum import IntEnum @@ -87,7 +87,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION @@ -1748,10 +1748,11 @@ class EngineCoreProc(EngineCore): encoder = MsgpackEncoder() # Send buffers to reuse. reuse_buffers: list[bytearray] = [] - # Keep references to outputs and buffers until zmq is finished - # with them (outputs may contain tensors/np arrays whose - # backing buffers were extracted for zero-copy send). - pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]() + # Payload buffers that can't be reused yet because zmq may still be + # sending them. + # Buffers of the zero-copy tensor/ndarray frames don't need tracking + # here: zmq itself holds a reference to each until it's done with it. + pending = deque[tuple[zmq.MessageTracker, bytearray]]() # We must set linger to ensure the ENGINE_CORE_DEAD # message is sent prior to closing the socket. @@ -1792,20 +1793,38 @@ class EngineCoreProc(EngineCore): # Reclaim buffers that zmq is finished with. while pending and pending[-1][0].done: - reuse_buffers.append(pending.pop()[2]) + reclaimed = pending.pop()[1] + if len(reuse_buffers) < max_reuse_bufs: + reuse_buffers.append(reclaimed) buffer = reuse_buffers.pop() if reuse_buffers else bytearray() buffers = encoder.encode_into(outputs, buffer) - tracker = sockets[client_index].send_multipart( - buffers, copy=False, track=True + tracker = self._send_msg_tracking_payload( + sockets[client_index], buffers ) if not tracker.done: - ref = outputs if len(buffers) > 1 else None - pending.appendleft((tracker, ref, buffer)) + pending.appendleft((tracker, buffer)) elif len(reuse_buffers) < max_reuse_bufs: # Limit the number of buffers to reuse. reuse_buffers.append(buffer) + @staticmethod + def _send_msg_tracking_payload( + socket: zmq.Socket, buffers: Sequence[bytestr] + ) -> zmq.MessageTracker: + """Send `buffers` as a zero-copy multipart message, returning a tracker + for the *first* frame. + + Used instead of `Socket.send_multipart()` because we reuse the buffer + passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a + tracker for the last frame only. + """ + more_flag = zmq.SNDMORE if len(buffers) > 1 else 0 + tracker = socket.send(buffers[0], more_flag, copy=False, track=True) + if more_flag: + socket.send_multipart(buffers[1:], copy=False) + return tracker + def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None: """Log and return a request-scoped error response for exceptions raised from the add request preprocessing in the input socket processing thread. diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 9460fdf48f7..febaa10ce61 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -7,7 +7,7 @@ import sys import uuid import weakref from abc import ABC, abstractmethod -from collections import Counter, defaultdict, deque +from collections import Counter, defaultdict from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass @@ -671,11 +671,6 @@ class MPClient(EngineCoreClient): self.core_engine: EngineIdentity = self.core_engines[0] self.utility_results: dict[int, AnyFuture] = {} - # Request objects which may contain pytorch-allocated tensors - # that we need to keep references to until zmq is done with the - # underlying data. - self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]() - # Start monitoring engine core processes for unexpected failures self.start_engine_core_monitor() @@ -707,14 +702,6 @@ class MPClient(EngineCoreClient): if self.resources.engine_dead: raise EngineDeadError() - def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any): - if not tracker.done: - self.pending_messages.appendleft((tracker, msg)) - - def free_pending_messages(self): - while self.pending_messages and self.pending_messages[-1][0].done: - self.pending_messages.pop() - def dp_engines_running(self) -> bool: return self.engines_running @@ -896,17 +883,12 @@ class SyncMPClient(MPClient): def _send_input(self, request_type: EngineCoreRequestType, request: Any): self.ensure_alive() - self.free_pending_messages() # (Identity, RequestType, SerializedRequest) msg = (self.core_engine, request_type.value, *self.encoder.encode(request)) - - if len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - self.input_socket.send_multipart(msg, copy=False) - return - - tracker = self.input_socket.send_multipart(msg, copy=False, track=True) - self.add_pending_message(tracker, request) + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + self.input_socket.send_multipart(msg, copy=False) def call_utility(self, method: str, *args) -> Any: call_id = uuid.uuid1().int >> 64 @@ -1129,32 +1111,16 @@ class AsyncMPClient(MPClient): engine = self.core_engine message = (request_type.value, *self.encoder.encode(request)) - return self._send_input_message(message, engine, request) + return self._send_input_message(message, engine) def _send_input_message( - self, message: tuple[bytestr, ...], engine: EngineIdentity, objects: Any + self, message: tuple[bytestr, ...], engine: EngineIdentity ) -> Awaitable[Any]: - """ - objects is a reference to retain until zmq is finished with the - buffers, in case they were extracted from tensors in the request. - """ self.ensure_alive() - self.free_pending_messages() - - msg = (engine,) + message - if not objects or len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - return self.input_socket.send_multipart(msg, copy=False) - - future: asyncio.Future[zmq.MessageTracker] - future = self.input_socket.send_multipart(msg, copy=False, track=True) - - def add_pending(f: asyncio.Future[zmq.MessageTracker]): - with contextlib.suppress(BaseException): - self.add_pending_message(f.result(), objects) - - future.add_done_callback(add_pending) - return future + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + return self.input_socket.send_multipart((engine,) + message, copy=False) async def call_utility_async(self, method: str, *args) -> Any: return await self._call_utility_async(method, *args, engine=self.core_engine) @@ -1169,7 +1135,7 @@ class AsyncMPClient(MPClient): EngineCoreRequestType.UTILITY.value, *self.encoder.encode((self.client_index, call_id, method, args)), ) - await self._send_input_message(message, engine, args) + await self._send_input_message(message, engine) self._ensure_output_queue_task() return await future From ba702e978e3bc6af3a601cee10fefdeb49e7e8b5 Mon Sep 17 00:00:00 2001 From: Yiliu Dong <1098822169@qq.com> Date: Wed, 29 Jul 2026 00:17:22 +0800 Subject: [PATCH 34/67] [Attention] Skip sparse indexer scoring for dense short prefills (#48407) Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .../layers/test_mla_short_prefill_indexer.py | 165 ++++++++++++++++++ .../layers/attention/mla_attention.py | 13 +- .../layers/attention/sparse_mla_attention.py | 4 + vllm/model_executor/layers/mla.py | 23 ++- .../layers/sparse_attn_indexer.py | 28 ++- vllm/model_executor/models/deepseek_v2.py | 3 + vllm/models/deepseek_v32/nvidia/attention.py | 6 +- 7 files changed, 230 insertions(+), 12 deletions(-) create mode 100644 tests/model_executor/layers/test_mla_short_prefill_indexer.py diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py new file mode 100644 index 00000000000..6e1e10e8b45 --- /dev/null +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer +from vllm.config import CUDAGraphMode +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata + +INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache" +MLA_LAYER = "model.layers.0.self_attn.attn" + + +def make_indexer_metadata( + *, + num_decodes: int = 0, + num_decode_tokens: int = 0, + num_prefills: int = 1, + num_prefill_tokens: int = 1, + slot_mapping: torch.Tensor | None = None, +) -> DeepseekV32IndexerMetadata: + if slot_mapping is None: + slot_mapping = torch.zeros(num_prefill_tokens, dtype=torch.long) + return DeepseekV32IndexerMetadata( + seq_lens=torch.empty(0, dtype=torch.int32), + max_seq_len=2048, + slot_mapping=slot_mapping, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + prefill=SimpleNamespace(chunks=[]) if num_prefills else None, + ) + + +def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0): + return SimpleNamespace( + num_decode_tokens=num_decode_tokens, + prefill=SimpleNamespace(use_dense_mha=use_dense_mha), + ) + + +@pytest.mark.parametrize( + "batch_kind", + ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"], +) +def test_short_prefill_updates_k_cache_before_scoring_decision( + monkeypatch: pytest.MonkeyPatch, + batch_kind: str, +): + slot_mapping = torch.tensor([63, 64, 127, 128, -1]) + mla_num_decode_tokens = 1 if batch_kind == "mla_decode" else 0 + runtime_mode = ( + CUDAGraphMode.FULL if batch_kind == "full" else CUDAGraphMode.PIECEWISE + ) + should_skip = batch_kind in ("short", "threshold_mismatch") + num_decodes = int(batch_kind == "threshold_mismatch") + num_decode_tokens = 3 if batch_kind == "threshold_mismatch" else 0 + num_prefills = 0 if batch_kind == "threshold_mismatch" else 2 + num_prefill_tokens = 0 if batch_kind == "threshold_mismatch" else 5 + if batch_kind == "threshold_mismatch": + # With MTP=3 the indexer threshold is four. A main MLA backend whose + # threshold is one (for example FlashMLA under DCP) still routes this + # three-token extend through dense prefill attention. + slot_mapping = slot_mapping[:3] + indexer_metadata = make_indexer_metadata( + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + slot_mapping=slot_mapping, + ) + if indexer_metadata.num_decodes: + indexer_metadata.decode = object() + mla_metadata = make_mla_metadata( + use_dense_mha=batch_kind != "force_mqa", + num_decode_tokens=mla_num_decode_tokens, + ) + + observed: dict[str, object] = {} + + monkeypatch.setattr( + sparse_indexer, + "get_forward_context", + lambda: SimpleNamespace( + attn_metadata={ + INDEXER_LAYER: indexer_metadata, + MLA_LAYER: mla_metadata, + }, + cudagraph_runtime_mode=runtime_mode, + ), + ) + monkeypatch.setattr( + sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16 + ) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: batch_kind == "capture", + ) + + def record_cache_update(k, kv_cache, slots, block_size, scale_fmt): + observed.update(k=k.clone(), slots=slots) + + monkeypatch.setattr( + sparse_indexer.ops, "indexer_k_quant_and_cache", record_cache_update + ) + + class ScoringReached(Exception): + pass + + def scoring_trigger(): + if should_skip: + pytest.fail("short dense-MHA prefill must not enter indexer scoring") + raise ScoringReached + + def scoring_decode(*args): + raise ScoringReached + + monkeypatch.setattr(sparse_indexer, "current_workspace_manager", scoring_trigger) + monkeypatch.setattr( + sparse_indexer, + "kv_cache_as_quant_view", + scoring_decode, + ) + + hidden_states = torch.full((7, 1), float("inf")) + k = torch.arange(28, dtype=torch.float32).reshape(7, 4) + topk_indices = torch.full((7, 2048), 17, dtype=torch.int32) + + def run_indexer(): + return sparse_indexer.sparse_attn_indexer( + hidden_states, + INDEXER_LAYER, + torch.empty(1), + torch.full((7, 1), float("inf")), + None, + k, + torch.full((7, 1), float("inf")), + 128, + "ue8m0", + 2048, + 4, + 4096, + 4096, + topk_indices, + False, + False, + MLA_LAYER, + ) + + if should_skip: + assert run_indexer() is topk_indices + assert torch.all(topk_indices == 17) + else: + with pytest.raises(ScoringReached): + run_indexer() + assert torch.all(topk_indices == -1) + + # K cache is always updated before the scoring decision. + torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()]) + assert observed["slots"] is slot_mapping diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1dea276d2d2..16fb961e563 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -769,13 +769,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): num_mha_tokens = q.size(0) - num_mqa_tokens if self.impl.is_sparse and num_mha_tokens > 0: - prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined] - use_mha = ( - self.prefill_backend is not None - and prefill_max_seq_len <= attn_metadata.topk_tokens # type: ignore[attr-defined] - and not self._vllm_config.attention_config.sparse_mla_force_mqa - ) - if not use_mha: + prefill_metadata = getattr(attn_metadata, "prefill", None) + if not getattr(prefill_metadata, "use_dense_mha", False): num_mqa_tokens = q.size(0) num_mha_tokens = 0 @@ -1409,6 +1404,10 @@ class MLACommonPrefillMetadata: q_data_type: torch.dtype | None = None output_dtype: torch.dtype | None = None prefill_backend: MLAPrefillBackend | None = None + # Whether the prefill suffix is routed through dense MHA. + # Indexer scoring may be skipped only for a pure-prefill batch, + # since decode tokens still consume top-k indices. + use_dense_mha: bool = False @dataclass diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 19cad7986bf..1463f9fb35b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -203,6 +203,10 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, + use_dense_mha=( + prefill_max_seq_len <= self.topk_tokens + and not self.vllm_config.attention_config.sparse_mla_force_mqa + ), ) self._prefill_backend.prepare_metadata(prefill) diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index ca4a20874a0..ac956461415 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -8,6 +8,7 @@ from vllm.config import CacheConfig from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention import MLAAttention from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.platforms import current_platform @dataclass @@ -65,6 +66,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): quant_config: QuantizationConfig | None = None, prefix: str = "", skip_topk: bool = False, + allow_short_prefill_indexer_scoring_skip: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -119,7 +121,26 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, ) - + indexer_op = getattr(self.indexer, "indexer_op", None) + if indexer_op is not None and hasattr( + indexer_op, "dense_mha_metadata_layer_name" + ): + enable_short_prefill_scoring_skip = ( + allow_short_prefill_indexer_scoring_skip + and not self.skip_topk + and not getattr(indexer_op, "use_pcp", False) + and current_platform.is_cuda() + ) + # The indexer and main MLA use independent decode thresholds and + # may classify the same short extend differently. Bind the main + # MLA layer name so the eager indexer op can check whether the + # batch's top-k indices will be consumed. + # PCP is excluded because indexer cache/scoring ownership differs + # across ranks and the no-consumer invariant has not been + # established there. + indexer_op.dense_mha_metadata_layer_name = ( + self.mla_attn.layer_name if enable_short_prefill_scoring_skip else "" + ) self.prefix = prefix def forward( diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 5b8e2bf008e..9a671be5563 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import get_current_vllm_config +from vllm.config import CUDAGraphMode, get_current_vllm_config from vllm.distributed import get_dcp_group, get_pcp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -310,6 +310,7 @@ def sparse_attn_indexer( topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -317,7 +318,8 @@ def sparse_attn_indexer( skip_topk_buffer_clear: bool = False, ) -> torch.Tensor: # careful! this will be None in dummy run - attn_metadata = get_forward_context().attn_metadata + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) @@ -357,6 +359,7 @@ def sparse_attn_indexer( topk_indices_buffer, skip_k_cache_insert, use_pcp, + dense_mha_metadata_layer_name, use_fp4_cache, ) attn_metadata_narrowed = attn_metadata[k_cache_prefix] @@ -402,6 +405,24 @@ def sparse_attn_indexer( scale_fmt, ) + # The indexer and main MLA may classify the same short extend differently + # because they use independent decode thresholds. Only the main MLA route + # can determine whether the top-k indices will be consumed. + if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL: + dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name) + if dense_mha_layer: + mla_metadata = attn_metadata.get(dense_mha_layer) + prefill_metadata = getattr(mla_metadata, "prefill", None) + if ( + getattr(prefill_metadata, "use_dense_mha", False) + and getattr(mla_metadata, "num_decode_tokens", -1) == 0 + and not torch.cuda.is_current_stream_capturing() + ): + # Deliberately leave the buffer untouched. Dense MHA does not + # consume top-k indices for this batch; clearing it would be + # unnecessary work. + return topk_indices_buffer + # The buffer must be pre-filled with -1 (the "no token" sentinel) before the # top-k kernels scatter valid indices into it. On the fused deepseek_v32 # nvidia path, _fused_norm_rope_kernel already cleared the same @@ -684,6 +705,7 @@ def sparse_attn_indexer_fake( topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, use_pcp: bool, + dense_mha_metadata_layer_name: LayerNameType, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -739,6 +761,7 @@ class SparseAttnIndexer(CustomOp): self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache + self.dense_mha_metadata_layer_name = "" # DCP scalars are constant for the run; resolve them here (config is set # during model construction) and pass them into the custom op, rather # than threading them through per-step metadata. @@ -800,6 +823,7 @@ class SparseAttnIndexer(CustomOp): self.topk_indices_buffer, self.skip_k_cache_insert, self.use_pcp, + _encode_layer_name(self.dense_mha_metadata_layer_name), self.use_fp4_cache, self.dcp_rank, self.dcp_world_size, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 4b92e351caa..bf67e040a15 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1178,6 +1178,9 @@ class DeepseekV2MLAAttention(nn.Module): # the V1 proposer. A frozen True would leave the draft reading a # never-written topk buffer. skip_topk=_skip_topk and not is_mtp_layer, + # Do not skip scoring for MTP layers: their top-k buffer may be + # reused by later draft iterations through index sharing. + allow_short_prefill_indexer_scoring_skip=not is_mtp_layer, ) def forward( diff --git a/vllm/models/deepseek_v32/nvidia/attention.py b/vllm/models/deepseek_v32/nvidia/attention.py index dcf955ad59b..0e604da87ba 100644 --- a/vllm/models/deepseek_v32/nvidia/attention.py +++ b/vllm/models/deepseek_v32/nvidia/attention.py @@ -494,8 +494,10 @@ class DeepseekV32Attention(MLAAttention): self.indexer.max_model_len, self.indexer.max_total_seq_len, self.topk_indices_buffer, - True, # skip_k_cache_insert - False, # use_fp4_cache + skip_k_cache_insert=True, + use_pcp=False, + dense_mha_metadata_layer_name="", + use_fp4_cache=False, # fused_norm_rope already cleared the topk buffer this forward. skip_topk_buffer_clear=True, ) From 01661cc57f48ce95c639efce7c88e6dd37349007 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 29 Jul 2026 00:21:12 +0800 Subject: [PATCH 35/67] [Rust][Benchmark] Make `vllm bench serve` Rust delegation opt-in (#50081) Signed-off-by: Bugen Zhao --- .../entrypoints/openai/test_dp_supervisor.py | 1 + tests/test_envs.py | 9 +++ vllm/entrypoints/cli/benchmark/main.py | 20 ++++++ vllm/entrypoints/cli/benchmark/serve.py | 62 +------------------ vllm/entrypoints/cli/main.py | 2 + vllm/entrypoints/cli/serve.py | 14 +++-- vllm/entrypoints/openai/dp_supervisor.py | 2 +- vllm/envs.py | 27 ++++---- 8 files changed, 61 insertions(+), 76 deletions(-) diff --git a/tests/entrypoints/openai/test_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 576e7ef16df..df80053deb5 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -201,6 +201,7 @@ def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch): monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None) monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None) monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None) + monkeypatch.setattr(dp_sup.envs, "VLLM_USE_RUST_FRONTEND", True, raising=False) monkeypatch.setattr( dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", diff --git a/tests/test_envs.py b/tests/test_envs.py index 56c04dd6f2e..5917c28fab2 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -145,6 +145,15 @@ def test_precompiled_install_flags_are_orthogonal() -> None: assert environment_variables["VLLM_USE_PRECOMPILED_RUST"]() is True +def test_rust_bench_auto_path_missing_fails_fast() -> None: + with ( + patch.dict(os.environ, {"VLLM_USE_RUST_BENCH": "1"}, clear=True), + patch("vllm.envs.os.path.isfile", return_value=False), + pytest.raises(FileNotFoundError, match="vllm-rs binary was not found"), + ): + environment_variables["VLLM_RUST_FRONTEND_PATH"]() + + class TestEnvWithChoices: """Test cases for env_with_choices function.""" diff --git a/vllm/entrypoints/cli/benchmark/main.py b/vllm/entrypoints/cli/benchmark/main.py index 1afac64b148..9ea49987091 100644 --- a/vllm/entrypoints/cli/benchmark/main.py +++ b/vllm/entrypoints/cli/benchmark/main.py @@ -2,18 +2,38 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import os import sys import typing +from vllm import envs from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.logger import init_logger if typing.TYPE_CHECKING: from vllm.utils.argparse_utils import FlexibleArgumentParser else: FlexibleArgumentParser = argparse.ArgumentParser +logger = init_logger(__name__) + + +def maybe_exec_rust_bench() -> None: + if sys.argv[1:3] != ["bench", "serve"] or not envs.VLLM_USE_RUST_BENCH: + return + + rust_cli = envs.VLLM_RUST_FRONTEND_PATH + if rust_cli is None: + raise RuntimeError( + "VLLM_USE_RUST_BENCH=1 requires VLLM_RUST_FRONTEND_PATH " + "to resolve to the vllm-rs binary." + ) + + logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli) + os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]]) + def _import_bench_subcommand_modules() -> None: # Imported lazily so `BenchmarkSubcommandBase` subclasses register only diff --git a/vllm/entrypoints/cli/benchmark/serve.py b/vllm/entrypoints/cli/benchmark/serve.py index 41a65273ba8..188afd6c703 100644 --- a/vllm/entrypoints/cli/benchmark/serve.py +++ b/vllm/entrypoints/cli/benchmark/serve.py @@ -1,68 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse -import os -import sys -from pathlib import Path -from vllm.benchmarks.serve import add_cli_args -from vllm.benchmarks.serve import main as python_main +from vllm.benchmarks.serve import add_cli_args, main from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase -from vllm.logger import init_logger from vllm.utils.argparse_utils import FlexibleArgumentParser -logger = init_logger(__name__) -_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs" -_RUST_SUPPORTED_DATASETS = frozenset( - { - "custom", - "hf", - "prefix_repetition", - "random", - "random-mm", - "random-rerank", - "sharegpt", - "sonnet", - "speed_bench", - } -) -_RUST_SUPPORTED_BACKENDS = frozenset( - { - "openai", - "openai-chat", - "openai-embeddings", - "openai-embeddings-chat", - "vllm", - "vllm-pooling", - "vllm-rerank", - } -) - - -def _rust_unsupported_reason(args: argparse.Namespace) -> str | None: - if args.dataset_name not in _RUST_SUPPORTED_DATASETS: - return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark" - if args.backend not in _RUST_SUPPORTED_BACKENDS: - return f"backend {args.backend!r} is not supported by the Rust benchmark" - return None - - -def _maybe_exec_rust_bench(args: argparse.Namespace) -> None: - if reason := _rust_unsupported_reason(args): - logger.info("Using Python benchmark: %s.", reason) - return - - if not _RUST_CLI_PATH.is_file(): - logger.warning( - "Rust benchmark binary not found at %s; falling back to Python.", - _RUST_CLI_PATH, - ) - return - - rust_cli = str(_RUST_CLI_PATH) - logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli) - os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]]) - class BenchmarkServingSubcommand(BenchmarkSubcommandBase): """The `serve` subcommand for `vllm bench`.""" @@ -76,5 +19,4 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase): @staticmethod def cmd(args: argparse.Namespace) -> None: - _maybe_exec_rust_bench(args) - python_main(args) + main(args) diff --git a/vllm/entrypoints/cli/main.py b/vllm/entrypoints/cli/main.py index fe0b339b3ed..3dc69dd3ad2 100644 --- a/vllm/entrypoints/cli/main.py +++ b/vllm/entrypoints/cli/main.py @@ -54,6 +54,8 @@ def main(): logger.info("Delegating entrypoint handling to vllm-omni") omni_main() else: + vllm.entrypoints.cli.benchmark.main.maybe_exec_rust_bench() + # For 'vllm bench *': use CPU instead of UnspecifiedPlatform by default if len(sys.argv) > 1 and sys.argv[1] == "bench": logger.debug( diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index d5e9b2bc874..08cb79f2081 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -58,6 +58,10 @@ class ServeSubcommand(CLISubcommand): uvloop.run(serve_grpc(args)) return + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) + if args.headless: if args.api_server_count is not None and args.api_server_count > 0: raise ValueError( @@ -103,7 +107,7 @@ class ServeSubcommand(CLISubcommand): # - Hybrid LB: Use local DP size (internal LB for local ranks only) # - Internal LB: Use full DP size if args.api_server_count is None: - if is_multi_port or is_external_lb or envs.VLLM_RUST_FRONTEND_PATH: + if is_multi_port or is_external_lb or rust_frontend_path: args.api_server_count = 1 elif is_hybrid_lb: args.api_server_count = args.data_parallel_size_local or 1 @@ -120,7 +124,7 @@ class ServeSubcommand(CLISubcommand): "Defaulting api_server_count to data_parallel_size (%d).", args.api_server_count, ) - elif envs.VLLM_RUST_FRONTEND_PATH and args.api_server_count > 1: + elif rust_frontend_path and args.api_server_count > 1: logger.warning( "Ignoring --api-server-count=%d when using rust front-end process", args.api_server_count, @@ -140,7 +144,7 @@ class ServeSubcommand(CLISubcommand): run_dp_supervisor(args) elif args.api_server_count < 1: run_headless(args) - elif args.api_server_count > 1 or envs.VLLM_RUST_FRONTEND_PATH: + elif args.api_server_count > 1 or rust_frontend_path: run_multi_api_server(args) else: # Single API server (this process). @@ -256,7 +260,9 @@ def run_headless(args: argparse.Namespace): def run_multi_api_server(args: argparse.Namespace): assert not args.headless - rust_frontend_path = envs.VLLM_RUST_FRONTEND_PATH + rust_frontend_path = ( + envs.VLLM_RUST_FRONTEND_PATH if envs.VLLM_USE_RUST_FRONTEND else None + ) num_api_servers: int = args.api_server_count assert num_api_servers > 0 diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index d669ec4d1d5..8ce6233c1ba 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -257,7 +257,7 @@ def _run_vllm_dp_server(child_args: argparse.Namespace) -> None: name = f"APIServer_DP{child_args.data_parallel_rank}" set_process_title(name) decorate_logs(name) - if envs.VLLM_RUST_FRONTEND_PATH: + if envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH: _run_rust_vllm_dp_server(child_args) else: _run_python_vllm_dp_server(child_args) diff --git a/vllm/envs.py b/vllm/envs.py index f7e01c33275..84ec5a8af85 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -153,6 +153,7 @@ if TYPE_CHECKING: K_SCALE_CONSTANT: int = 200 V_SCALE_CONSTANT: int = 100 VLLM_USE_RUST_FRONTEND: bool = False + VLLM_USE_RUST_BENCH: bool = False VLLM_RUST_FRONTEND_PATH: str | None = "auto" VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 @@ -548,22 +549,24 @@ def _deprecated_triton_attn_use_td() -> None: return None -def _resolve_rust_frontend_path() -> str | None: - """Resolve the Rust frontend binary path. +def _resolve_rust_cli_path() -> str | None: + """Resolve the vllm-rs binary path. - Returns None if VLLM_USE_RUST_FRONTEND is not enabled. + Returns None unless VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH is enabled. When enabled, resolves VLLM_RUST_FRONTEND_PATH ("auto" by default) to the actual binary path. """ - use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) + use_rust = bool(int(os.environ.get("VLLM_USE_RUST_FRONTEND", "0"))) or bool( + int(os.environ.get("VLLM_USE_RUST_BENCH", "0")) + ) raw = os.environ.get("VLLM_RUST_FRONTEND_PATH", "auto") if not use_rust: if os.environ.get("VLLM_RUST_FRONTEND_PATH") is not None: logger.warning( - "VLLM_RUST_FRONTEND_PATH is set but VLLM_USE_RUST_FRONTEND " - "is not enabled. The Rust frontend will not be used. " - "Set VLLM_USE_RUST_FRONTEND=1 to enable it." + "VLLM_RUST_FRONTEND_PATH is set without enabling " + "VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH. " + "Set one of them to 1 to use the vllm-rs binary." ) return None @@ -1340,10 +1343,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_RUST_FRONTEND": lambda: bool( int(os.getenv("VLLM_USE_RUST_FRONTEND", "0")) ), - # Path to the Rust frontend binary. Defaults to "auto" which discovers - # the binary installed with the vllm package. Only used when - # VLLM_USE_RUST_FRONTEND=1. - "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_frontend_path(), + # If set, use the packaged Rust client for `vllm bench serve`. + "VLLM_USE_RUST_BENCH": lambda: bool(int(os.getenv("VLLM_USE_RUST_BENCH", "0"))), + # Path to the vllm-rs binary. Defaults to "auto" which discovers the + # binary installed with the vllm package. Used when VLLM_USE_RUST_FRONTEND=1 + # or VLLM_USE_RUST_BENCH=1. + "VLLM_RUST_FRONTEND_PATH": lambda: _resolve_rust_cli_path(), # If set, vllm will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` From 4f56321d7ec25dd041c7bf0aa47b4eadd075a7e5 Mon Sep 17 00:00:00 2001 From: jiacao-amd Date: Tue, 28 Jul 2026 09:34:16 -0700 Subject: [PATCH 36/67] [ROCm] Cache fp32 upcast of static e8m0 weight scale in AITER scaled_mm (#47773) Signed-off-by: jiacao-amd Co-authored-by: TJian --- .../kernels/linear/scaled_mm/aiter.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 1b39491ab34..da8f69fa97b 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -9,6 +9,9 @@ from vllm._aiter_ops import ( rocm_aiter_ops, ) from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, ) @@ -16,6 +19,7 @@ from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from .BlockScaledMMLinearKernel import ( + FP8BlockParams, Fp8BlockScaledMMLinearKernel, ) from .cutlass import CutlassInt8ScaledMMLinearKernel @@ -375,6 +379,17 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + + params = FP8BlockParams.from_layer(layer) + if params.weight_scale_inv is not None: + ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV + else: + ws, attr = params.weight_scale, params.WEIGHT_SCALE + if ws is not None and ws.dtype == torch.float8_e8m0fnu: + replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous()) + @classmethod def is_supported(cls, compute_capability=None): return ( @@ -406,19 +421,12 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): Bs: torch.Tensor, ) -> torch.Tensor: if As.dtype != Bs.dtype: - from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, - ) - if As.dtype == torch.float8_e8m0fnu: As = _upcast_e8m0_to_fp32(As).contiguous() else: As = As.to(torch.float32) - if Bs.dtype == torch.float8_e8m0fnu: - Bs = _upcast_e8m0_to_fp32(Bs).contiguous() - else: - Bs = Bs.to(torch.float32) + Bs = Bs.to(torch.float32) out_dtype = self.config.out_dtype if self.use_triton: From 05a08148631ab2d852c1984ddb3abebad8039a29 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 28 Jul 2026 11:50:10 -0500 Subject: [PATCH 37/67] [ROCm] Fix and optimize GPT-J-style MRoPE (#49906) Signed-off-by: Andreas Karatzas --- tests/kernels/core/test_mrope.py | 10 +- .../layers/rotary_embedding/mrope.py | 149 +++++++++++------- 2 files changed, 100 insertions(+), 59 deletions(-) diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 29051b4a00c..6f64fabfe9b 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -38,6 +38,7 @@ def generate_test_data( class MRoPETestInfo(NamedTuple): model_name: str + is_neox_style: bool = True # https://github.com/pytorch/pytorch/blob/main/torch/testing/_comparison.py#L1317 atol: float = 1e-2 rtol: float = 1.6e-2 @@ -45,7 +46,10 @@ class MRoPETestInfo(NamedTuple): MODELS_TO_TEST = [ - MRoPETestInfo(model_name="zai-org/GLM-4.1V-9B-Thinking"), + MRoPETestInfo( + model_name="zai-org/GLM-4.1V-9B-Thinking", + is_neox_style=False, + ), MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"), MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"), @@ -92,7 +96,7 @@ def test_mrope( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings @@ -162,7 +166,7 @@ def test_mrope_torch_compile_tracing( if hasattr(config, "head_dim") else config.hidden_size // total_num_heads ) - is_neox_style = True + is_neox_style = model_info.is_neox_style max_position = config.max_position_embeddings mrope_helper_class = get_rope( diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 3c946dd130c..29ce9e5000d 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -5,6 +5,7 @@ import numpy as np import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .base import RotaryEmbeddingBase @@ -24,16 +25,17 @@ def _triton_mrope_forward( rd: tl.constexpr, pad_n_qh: tl.constexpr, pad_n_kh: tl.constexpr, - pad_hd: tl.constexpr, + pad_rd: tl.constexpr, mrope_section_t: tl.constexpr, mrope_section_h: tl.constexpr, mrope_section_w: tl.constexpr, is_interleaved: tl.constexpr, + is_neox_style: tl.constexpr, ): # Adapted from # https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/qwen2vl_mrope.py # This version supports flatten input tensors from vllm - # and supports cos and sin cache with shape (3, num_tokens, head_dim // 2) + # and supports cos and sin cache with shape (3, num_tokens, rotary_dim // 2) # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary pid = tl.program_id(0) # locate start address @@ -44,9 +46,9 @@ def _triton_mrope_forward( # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position # m of this program instance # #################################################################### - # Note: cos and sin now have shape (3, num_tokens, head_dim // 2) + # Note: cos and sin now have shape (3, num_tokens, rotary_dim // 2) - # Updated stride calculation for half head_dim + # Updated stride calculation for half rotary_dim half_rd = rd // 2 t_cos = cos + pid * half_rd h_cos = t_cos + num_tokens * half_rd @@ -55,12 +57,17 @@ def _triton_mrope_forward( h_sin = t_sin + num_tokens * half_rd w_sin = h_sin + num_tokens * half_rd - # Updated offsets for half head_dim - cos_offsets = tl.arange(0, pad_hd // 2) + # Updated offsets for half rotary_dim + cos_offsets = tl.arange(0, pad_rd // 2) if is_interleaved: - h_mask = ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) - w_mask = ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) - t_mask = ~(h_mask | w_mask) + valid_mask = cos_offsets < half_rd + h_mask = ( + valid_mask & ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h) + ) + w_mask = ( + valid_mask & ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w) + ) + t_mask = valid_mask & ~(h_mask | w_mask) else: t_end = mrope_section_t h_end = t_end + mrope_section_h @@ -79,55 +86,74 @@ def _triton_mrope_forward( sin_row = t_sin_row + h_sin_row + w_sin_row # #################################################################### - # Load the left and right half of q and k for the current - # program instance (i.e. for the current token) separately + # Load the two values in each rotary pair for the current token. + # NeoX pairs the first and second halves, while GPT-J pairs + # adjacent values. # #################################################################### - # left half of the head - first_half_q_offsets = ( - tl.arange(0, pad_n_qh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_half_k_offsets = ( - tl.arange(0, pad_n_kh)[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :] - ) - first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) - first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( - tl.arange(0, pad_hd // 2)[None, :] < rd // 2 - ) + if is_neox_style: + rotary_offsets = tl.arange(0, pad_rd // 2) + first_q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + first_k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + first_q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd // 2 + ) + first_k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd // 2 + ) - q_tile_1 = tl.load(q_ptr + first_half_q_offsets, mask=first_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_1 = tl.load(k_ptr + first_half_k_offsets, mask=first_k_mask, other=0).to( - sin_row.dtype - ) + q_tile_1 = tl.load(q_ptr + first_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_1 = tl.load(k_ptr + first_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - # right half of the head - second_half_q_offsets = first_half_q_offsets + (rd // 2) - second_half_k_offsets = first_half_k_offsets + (rd // 2) - second_q_mask = first_q_mask - second_k_mask = first_k_mask + second_q_offsets = first_q_offsets + (rd // 2) + second_k_offsets = first_k_offsets + (rd // 2) + q_tile_2 = tl.load(q_ptr + second_q_offsets, mask=first_q_mask, other=0).to( + sin_row.dtype + ) + k_tile_2 = tl.load(k_ptr + second_k_offsets, mask=first_k_mask, other=0).to( + sin_row.dtype + ) - q_tile_2 = tl.load(q_ptr + second_half_q_offsets, mask=second_q_mask, other=0).to( - sin_row.dtype - ) - k_tile_2 = tl.load(k_ptr + second_half_k_offsets, mask=second_k_mask, other=0).to( - sin_row.dtype - ) + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + tl.store(q_ptr + first_q_offsets, new_q_tile_1, mask=first_q_mask) + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + tl.store(q_ptr + second_q_offsets, new_q_tile_2, mask=first_q_mask) - # y = [x1, x2] * [cos, cos] + [-x2, x1] * [sin, sin] - # Since cos and sin are now half-size, - # we use the same cos_row and sin_row for both halves - new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row - tl.store(q_ptr + first_half_q_offsets, new_q_tile_1, mask=first_q_mask) - new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row - tl.store(q_ptr + second_half_q_offsets, new_q_tile_2, mask=second_q_mask) + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + tl.store(k_ptr + first_k_offsets, new_k_tile_1, mask=first_k_mask) + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + tl.store(k_ptr + second_k_offsets, new_k_tile_2, mask=first_k_mask) + else: + # Load and store adjacent rotary pairs contiguously. Using stride-two + # even/odd offsets makes Triton emit scalar 16-bit memory operations on + # AMD, while split/interleave only rearranges values in registers. + rotary_offsets = tl.arange(0, pad_rd) + q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :] + k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :] + q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & ( + rotary_offsets[None, :] < rd + ) + k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & ( + rotary_offsets[None, :] < rd + ) - new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row - tl.store(k_ptr + first_half_k_offsets, new_k_tile_1, mask=first_k_mask) - new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row - tl.store(k_ptr + second_half_k_offsets, new_k_tile_2, mask=second_k_mask) + q_tile = tl.load(q_ptr + q_offsets, mask=q_mask, other=0).to(sin_row.dtype) + k_tile = tl.load(k_ptr + k_offsets, mask=k_mask, other=0).to(sin_row.dtype) + q_tile_1, q_tile_2 = tl.split(tl.reshape(q_tile, (pad_n_qh, pad_rd // 2, 2))) + k_tile_1, k_tile_2 = tl.split(tl.reshape(k_tile, (pad_n_kh, pad_rd // 2, 2))) + + new_q_tile_1 = q_tile_1 * cos_row - q_tile_2 * sin_row + new_q_tile_2 = q_tile_2 * cos_row + q_tile_1 * sin_row + new_q_tile = tl.interleave(new_q_tile_1, new_q_tile_2) + tl.store(q_ptr + q_offsets, new_q_tile, mask=q_mask) + + new_k_tile_1 = k_tile_1 * cos_row - k_tile_2 * sin_row + new_k_tile_2 = k_tile_2 * cos_row + k_tile_1 * sin_row + new_k_tile = tl.interleave(new_k_tile_1, new_k_tile_2) + tl.store(k_ptr + k_offsets, new_k_tile, mask=k_mask) def triton_mrope( @@ -139,23 +165,26 @@ def triton_mrope( head_size: int, rotary_dim: int, mrope_interleaved: bool, + is_neox_style: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Qwen2VL mrope kernel. Args: q: [num_tokens, num_heads * head_size] k: [num_tokens, num_kv_heads * head_size] - cos: [3, num_tokens, head_size //2 ] + cos: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) - sin: [3, num_tokens, head_size //2 ] + sin: [3, num_tokens, rotary_dim // 2] (T/H/W positions with multimodal inputs) mrope_section: [t, h, w] head_size: int + is_neox_style: Whether rotary pairs use split-half (NeoX) or + adjacent (GPT-J) layout. """ n_row, n_q_head_head_dim = q.shape n_q_head = n_q_head_head_dim // head_size n_kv_head = k.shape[1] // head_size - pad_hd = triton.next_power_of_2(head_size) + pad_rd = triton.next_power_of_2(rotary_dim) pad_n_q_head = triton.next_power_of_2(n_q_head) pad_n_kv_head = triton.next_power_of_2(n_kv_head) @@ -166,6 +195,11 @@ def triton_mrope( cos = cos.contiguous() sin = sin.contiguous() + # Small adjacent-pair tiles perform best with one wave per program on + # ROCm. Keep the existing launch shape for larger rotary dimensions, + # NeoX, and other backends. + use_single_wave = current_platform.is_rocm() and not is_neox_style and pad_rd <= 64 + num_warps = 1 if use_single_wave else 4 _triton_mrope_forward[(n_row,)]( q, k, @@ -178,11 +212,13 @@ def triton_mrope( rotary_dim, pad_n_q_head, pad_n_kv_head, - pad_hd, + pad_rd, mrope_section[0], mrope_section[1], mrope_section[2], mrope_interleaved, + is_neox_style, + num_warps=num_warps, ) return q, k @@ -349,6 +385,7 @@ class MRotaryEmbedding(RotaryEmbeddingBase): self.head_size, self.rotary_dim, self.mrope_interleaved, + self.is_neox_style, ) return q.reshape(query_shape), k.reshape(key_shape) From 8a7b3c299053efbca2669081ddf50a46b6a9d149 Mon Sep 17 00:00:00 2001 From: Brian Dellabetta Date: Tue, 28 Jul 2026 13:03:08 -0400 Subject: [PATCH 38/67] [compressed-tensors] update `find_matched_target` order to prioritize fused name matches over class match (#49483) Signed-off-by: Brian Dellabetta --- .../layers/quantization/compressed_tensors/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py index afb899cd6d7..872771af9ab 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py @@ -145,8 +145,8 @@ def find_matched_target( matched_target = ( _find_first_match(layer_name, targets) - or _find_first_match(module.__class__.__name__, targets, True) or _match_fused_layer(layer_name, targets, fused_mapping) + or _find_first_match(module.__class__.__name__, targets, True) ) return matched_target From 1db989bbf14b9896c1306fd722ced9a9b1238465 Mon Sep 17 00:00:00 2001 From: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:57:31 +0800 Subject: [PATCH 39/67] [Bugfix][Multimodal] Fix video temporal padding estimates (#49030) Signed-off-by: xiaoming <1259730330@qq.com> --- .../multimodal/processing/test_glm4_1v.py | 25 +++++++++++++++++++ vllm/model_executor/models/glm4_1v.py | 4 +-- vllm/model_executor/models/kanana_v.py | 4 +-- vllm/model_executor/models/keye.py | 2 +- .../model_executor/models/llava_onevision2.py | 2 +- vllm/model_executor/models/mimo_v2_omni.py | 2 +- vllm/model_executor/models/qwen2_vl.py | 4 +-- 7 files changed, 34 insertions(+), 9 deletions(-) diff --git a/tests/models/multimodal/processing/test_glm4_1v.py b/tests/models/multimodal/processing/test_glm4_1v.py index 8a777832826..e45e741c5b4 100644 --- a/tests/models/multimodal/processing/test_glm4_1v.py +++ b/tests/models/multimodal/processing/test_glm4_1v.py @@ -63,6 +63,31 @@ def test_encoder_cudagraph_uses_model_video_frame_limit(): assert Glm4vForConditionalGeneration.get_max_frames_per_video(model) == 600 +@pytest.mark.parametrize( + ("temporal_patch_size", "expected_grid_t"), + [(2, 9), (4, 5), (8, 3)], +) +def test_vision_info_rounds_up_temporal_frames( + temporal_patch_size: int, + expected_grid_t: int, +): + info = Mock(spec=Glm4vProcessingInfo) + vision_config = info.get_hf_config.return_value.vision_config + vision_config.patch_size = 14 + vision_config.spatial_merge_size = 2 + vision_config.temporal_patch_size = temporal_patch_size + + _, num_vision_tokens = Glm4vProcessingInfo._get_vision_info( + info, + image_width=28, + image_height=28, + num_frames=17, + do_resize=False, + ) + + assert num_vision_tokens == expected_grid_t + + @pytest.mark.parametrize("model_id", ["zai-org/GLM-4.1V-9B-Thinking"]) @pytest.mark.parametrize("expected_toks_per_frame", [299]) @pytest.mark.parametrize( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 787e53e8df2..19389d84463 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -1082,8 +1082,8 @@ class Glm4vProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/kanana_v.py b/vllm/model_executor/models/kanana_v.py index 125d7e71c7b..b1a5f78b1b3 100644 --- a/vllm/model_executor/models/kanana_v.py +++ b/vllm/model_executor/models/kanana_v.py @@ -409,8 +409,8 @@ class KananaVProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index dd1fb892ad1..c3d69836a79 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -983,7 +983,7 @@ class KeyeProcessingInfo(BaseProcessingInfo): else: preprocessed_size = ImageSize(width=image_width, height=image_height) - padded_num_frames = num_frames + num_frames % temporal_patch_size + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size diff --git a/vllm/model_executor/models/llava_onevision2.py b/vllm/model_executor/models/llava_onevision2.py index 58179ec00de..9dc11d493a0 100644 --- a/vllm/model_executor/models/llava_onevision2.py +++ b/vllm/model_executor/models/llava_onevision2.py @@ -1345,7 +1345,7 @@ class LlavaOnevision2ProcessingInfo(BaseProcessingInfo): preprocessed = ImageSize(width=rw, height=rh) else: preprocessed = ImageSize(width=image_width, height=image_height) - padded_frames = num_frames + num_frames % temporal_patch_size + padded_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_frames // temporal_patch_size, 1) grid_h = preprocessed.height // patch_size grid_w = preprocessed.width // patch_size diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index d0d9589ae1d..747cb0e88b2 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -715,7 +715,7 @@ class MiMoV2OmniProcessingInfo(BaseProcessingInfo): effective_frames = num_frames * tokens_per_second else: effective_frames = num_frames - padded_num_frames = effective_frames + effective_frames % temporal_patch_size + padded_num_frames = effective_frames + (-effective_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size grid_w = preprocessed_size.width // patch_size diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 539f141cbaa..e2e9f245248 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -898,8 +898,8 @@ class Qwen2VLProcessingInfo(BaseProcessingInfo): preprocessed_size = ImageSize(width=image_width, height=image_height) # NOTE: Frames are padded to be divisible by `temporal_patch_size` - # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294 - padded_num_frames = num_frames + num_frames % temporal_patch_size + # https://github.com/huggingface/transformers/blob/v5.13.0/src/transformers/models/qwen2_vl/video_processing_qwen2_vl.py#L249-L252 + padded_num_frames = num_frames + (-num_frames % temporal_patch_size) grid_t = max(padded_num_frames // temporal_patch_size, 1) grid_h = preprocessed_size.height // patch_size From 6c7e679f048dc6123caecc3985150766e455ff22 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Wed, 29 Jul 2026 02:08:50 +0800 Subject: [PATCH 40/67] [ROCm][Bugfix] Sanitize AITER paged-MQA logits before sparse top-k for DeepSeek-V4 (#49714) Signed-off-by: shen-shanshan <467638484@qq.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 64 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 1 + 2 files changed, 65 insertions(+) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 77e068ab171..6fe2a3e7758 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -196,6 +198,68 @@ def _ragged_from_rows( ) +@torch.inference_mode() +def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + + class FakeWorkspaceManager: + def get_simultaneous(self, *shapes_and_dtypes): + return [ + torch.empty(shape, dtype=dtype, device=device) + for shape, dtype in shapes_and_dtypes + ] + + def fake_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_logits, + context_lens, + block_tables, + max_seq_len, + **kwargs, + ): + del ( + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_seq_len, + kwargs, + ) + out_logits.fill_(float("nan")) + + monkeypatch.setattr(mod, "_ON_GFX942", False) + monkeypatch.setattr(mod, "_ON_GFX950", True) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + mod, + "paged_mqa_logits_module", + lambda: SimpleNamespace(deepgemm_fp8_paged_mqa_logits=fake_paged_mqa_logits), + ) + monkeypatch.setattr( + mod, "current_workspace_manager", lambda: FakeWorkspaceManager() + ) + + q_fp8 = torch.empty((1, 1, 1, 1), dtype=torch.uint8, device=device) + kv_cache_fp8 = torch.empty((1, 1, 1, 5), dtype=torch.uint8, device=device) + logits = mod.rocm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + torch.empty((1, 1), dtype=torch.float32, device=device), + torch.ones(1, dtype=torch.int32, device=device), + torch.zeros((1, 1), dtype=torch.int32, device=device), + torch.empty(0, dtype=torch.int32, device=device), + 1, + ) + + assert not torch.isnan(logits).any() + + @torch.inference_mode() def test_compute_global_topk_ragged_indices_and_indptr() -> None: from vllm.models.deepseek_v4.amd.rocm import ( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 6a28324208f..63bad3cfad0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -442,6 +442,7 @@ def rocm_fp8_paged_mqa_logits( KVBlockSize=block_size, WavePerEU=2, ) + out_logits.nan_to_num_(float("-inf")) return out_logits deepgemm_fp8_paged_mqa_logits_stage1 = ( aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1 From 118bcde449c923ff20fe6336a009e2cef2ba5423 Mon Sep 17 00:00:00 2001 From: johnnyychiu Date: Tue, 28 Jul 2026 20:00:06 +0100 Subject: [PATCH 41/67] [BugFix] Fix clang spinloop mwaitx include (#45532) Signed-off-by: johnny --- csrc/spinloop.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/spinloop.cpp b/csrc/spinloop.cpp index c29e48a5f0e..3285b9a3ded 100644 --- a/csrc/spinloop.cpp +++ b/csrc/spinloop.cpp @@ -7,7 +7,7 @@ extern "C" { #if defined(__i386__) || defined(__x86_64__) #include - #include + #include #endif #if defined(CLOCK_MONOTONIC_RAW) From d552a68645b60fe194ed488e9c3df211e4d2fa6f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 29 Jul 2026 03:31:15 +0800 Subject: [PATCH 42/67] [Rust Frontend] Extract shared tracing setup logic into `vllm-tracing` (#50129) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 14 +++++++++++--- rust/Cargo.toml | 2 ++ rust/src/bench/Cargo.toml | 2 +- rust/src/bench/src/main.rs | 12 +----------- rust/src/cmd/Cargo.toml | 3 +-- rust/src/cmd/src/main.rs | 3 +-- rust/src/tracing/Cargo.toml | 14 ++++++++++++++ .../src/{cmd/src/logging.rs => tracing/src/lib.rs} | 6 ++++-- 8 files changed, 35 insertions(+), 21 deletions(-) create mode 100644 rust/src/tracing/Cargo.toml rename rust/src/{cmd/src/logging.rs => tracing/src/lib.rs} (98%) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 5d2faecdb36..ee4d24622be 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5503,9 +5503,9 @@ dependencies = [ "tokio", "tokio-stream", "tracing", - "tracing-subscriber", "url", "uuid", + "vllm-tracing", ] [[package]] @@ -5570,17 +5570,16 @@ dependencies = [ "serde_json", "serde_with", "thiserror-ext", - "time", "tokio", "tokio-util", "tracing", - "tracing-subscriber", "uuid", "vllm-bench", "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", + "vllm-tracing", ] [[package]] @@ -5827,6 +5826,15 @@ dependencies = [ "vllm-parser", ] +[[package]] +name = "vllm-tracing" +version = "0.1.0" +dependencies = [ + "time", + "tracing", + "tracing-subscriber", +] + [[package]] name = "walkdir" version = "2.5.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index a33a13ff69b..42bda825cf8 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,6 +13,7 @@ members = [ "src/server", "src/text", "src/tokenizer", + "src/tracing", ] resolver = "3" @@ -143,6 +144,7 @@ vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } +vllm-tracing = { path = "src/tracing" } winnow = { version = "1.0.2", features = ["simd"] } xgrammar-structural-tag = "0.2.0" zeromq = { version = "0.6.0", default-features = false, features = [ diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 2da0f13f1ce..960e7a62f7d 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -32,9 +32,9 @@ tokenizers.workspace = true tokio.workspace = true tokio-stream.workspace = true tracing.workspace = true -tracing-subscriber.workspace = true url.workspace = true uuid.workspace = true +vllm-tracing.workspace = true [lints] workspace = true diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index b9ebfc2a0ac..1764983e448 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -19,18 +19,8 @@ struct Cli { args: vllm_bench::BenchServeArgs, } -// TODO: unify the tracing subscriber used by different binaries. -fn init_tracing() { - let filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); - let _ = tracing_subscriber::fmt() - .with_env_filter(filter) - .with_writer(std::io::stderr) - .try_init(); -} - fn main() -> anyhow::Result<()> { - init_tracing(); + vllm_tracing::init_tracing("Bench"); let cli = Cli::parse(); vllm_bench::prepare_process(); diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index a6955059c26..a326a0f9992 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -23,17 +23,16 @@ serde.workspace = true serde_json.workspace = true serde_with.workspace = true thiserror-ext.workspace = true -time.workspace = true tokio = { workspace = true, features = ["signal"] } tokio-util.workspace = true tracing.workspace = true -tracing-subscriber.workspace = true uuid.workspace = true vllm-bench.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true vllm-server.workspace = true +vllm-tracing.workspace = true [dev-dependencies] expect-test.workspace = true diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index 0806f7b75d4..86c9ab6f934 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project mod cli; -mod logging; use std::env; use std::ffi::OsStr; @@ -89,7 +88,7 @@ fn main() -> Result<()> { "serve" | "frontend" => "RustFrontend", _ => "Rust", }; - logging::init_tracing(process_label); + vllm_tracing::init_tracing(process_label); let cli = Cli::parse(); diff --git a/rust/src/tracing/Cargo.toml b/rust/src/tracing/Cargo.toml new file mode 100644 index 00000000000..e003720814e --- /dev/null +++ b/rust/src/tracing/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "vllm-tracing" +version.workspace = true +edition.workspace = true +description = "Shared tracing subscriber and log formatting for vLLM Rust binaries" +license.workspace = true + +[dependencies] +time.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/rust/src/cmd/src/logging.rs b/rust/src/tracing/src/lib.rs similarity index 98% rename from rust/src/cmd/src/logging.rs rename to rust/src/tracing/src/lib.rs index 936b3692442..eb6da9c0b24 100644 --- a/rust/src/cmd/src/logging.rs +++ b/rust/src/tracing/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +//! Shared tracing subscriber and log formatting for vLLM Rust binaries. + use std::{env, fmt, process}; use time::UtcOffset; @@ -26,8 +28,8 @@ const RESET: &str = "\x1b[0m"; const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] = format_description!("[month]-[day] [hour]:[minute]:[second]"); -/// Install the process-wide vLLM-style tracing subscriber for the CLI binary. -pub(crate) fn init_tracing(process_label: &str) { +/// Install the process-wide vLLM-style tracing subscriber. +pub fn init_tracing(process_label: &str) { let filter = build_targets_filter( env::var("VLLM_LOGGING_LEVEL").ok().as_deref(), env::var("RUST_LOG").ok().as_deref(), From 0b6aa3c47ce69a1f3c8a19cafe3b9dc2871d1f6b Mon Sep 17 00:00:00 2001 From: Siddhant Bharti <42143349+siddhant-bharti@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:34:22 -0700 Subject: [PATCH 43/67] [Bugfix][Spec Decode] Size DFlash query buffers for cudagraph-padded batches (#50065) Signed-off-by: siddhant-bharti Co-authored-by: Claude Fable 5 --- vllm/v1/spec_decode/dflash.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 5d1c0629218..626dd36dea9 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -38,8 +38,12 @@ class DFlashProposer(SpecDecodeBaseProposer): # Only next_token_ids and mask tokens are query tokens, all other context is K/V self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens) + self.max_padded_query_tokens = max( + self.max_query_tokens, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) # Positions covers both context states + query states - self.max_positions = self.max_num_tokens + self.max_query_tokens + self.max_positions = self.max_num_tokens + self.max_padded_query_tokens # Separate context buffers to keep query buffer addresses stable for CUDA graphs self._context_slot_mapping_buffer = torch.zeros( @@ -48,7 +52,7 @@ class DFlashProposer(SpecDecodeBaseProposer): device=device, ) self._slot_mapping_buffer = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, ) @@ -58,7 +62,7 @@ class DFlashProposer(SpecDecodeBaseProposer): device=device, ) self.positions = torch.zeros( - self.max_query_tokens, + self.max_padded_query_tokens, dtype=torch.int64, device=device, ) From 2899dca8432d40632987b0ec24253a8fe6df2710 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 29 Jul 2026 05:31:13 +0800 Subject: [PATCH 44/67] [Model] Add Kimi K3 support: Rust frontend [1/2] (#50104) Signed-off-by: Bugen Zhao --- pyproject.toml | 5 +- rust/Cargo.lock | 2 +- rust/Cargo.toml | 2 +- rust/src/chat/src/backend/hf.rs | 3 +- rust/src/chat/src/lib.rs | 13 +- rust/src/chat/src/parser/reasoning/mod.rs | 2 + rust/src/chat/src/parser/tool/mod.rs | 2 + rust/src/chat/src/parser/unified.rs | 26 +- .../chat/src/renderer/deepseek_v32/tests.rs | 2 +- .../chat/src/renderer/deepseek_v4/tests.rs | 2 +- rust/src/chat/src/renderer/harmony/tests.rs | 2 +- rust/src/chat/src/renderer/inkling/tests.rs | 2 +- .../src/chat/src/renderer/kimi_k3/encoding.rs | 608 +++++++++ .../fixtures/controls_thinking_off_input.json | 15 + .../fixtures/controls_thinking_off_output.txt | 2 + .../dynamic_system_tool_declare_input.json | 62 + .../dynamic_system_tool_declare_output.txt | 14 + .../history_preserve_and_image_input.json | 31 + .../history_preserve_and_image_output.txt | 2 + .../tools_history_and_required_input.json | 87 ++ .../tools_history_and_required_output.txt | 8 + rust/src/chat/src/renderer/kimi_k3/mod.rs | 39 + rust/src/chat/src/renderer/kimi_k3/tests.rs | 267 ++++ rust/src/chat/src/renderer/mod.rs | 2 + rust/src/chat/src/renderer/selection.rs | 9 +- rust/src/chat/src/renderer/test_utils.rs | 52 +- rust/src/chat/src/request.rs | 5 + rust/src/chat/tests/roundtrip.rs | 42 +- rust/src/cmd/src/cli/tests.rs | 2 +- rust/src/parser/src/unified/kimi_k3.rs | 1149 +++++++++++++++++ .../src/unified/kimi_k3/structural_tag.rs | 503 ++++++++ rust/src/parser/src/unified/mod.rs | 2 + .../routes/openai/chat_completions/convert.rs | 73 +- rust/src/server/src/routes/tokenize/types.rs | 1 + rust/src/tokenizer/src/tiktoken.rs | 4 +- 35 files changed, 3003 insertions(+), 39 deletions(-) create mode 100644 rust/src/chat/src/renderer/kimi_k3/encoding.rs create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/mod.rs create mode 100644 rust/src/chat/src/renderer/kimi_k3/tests.rs create mode 100644 rust/src/parser/src/unified/kimi_k3.rs create mode 100644 rust/src/parser/src/unified/kimi_k3/structural_tag.rs diff --git a/pyproject.toml b/pyproject.toml index 906702fa075..a4c7ff487a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,10 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", - "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs", + "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] + "rust/src/tokenizer/src/incremental.rs"] ignore-hidden = false [tool.typos.default] diff --git a/rust/Cargo.lock b/rust/Cargo.lock index ee4d24622be..38f1c257756 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5" +source = "git+https://github.com/smg-project/llm-multimodal?rev=15adba5e025d8636ba4a334fb379b1371f6196a1#15adba5e025d8636ba4a334fb379b1371f6196a1" dependencies = [ "anyhow", "base64 0.22.1", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 42bda825cf8..09f55cf07cd 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -59,7 +59,7 @@ indexmap = "2.13.0" indicatif = "0.18.4" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index d4999d07cd0..583a5fb936e 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -20,7 +20,7 @@ use crate::output::{ use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; use crate::renderer::{ DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, - InklingChatRenderer, + InklingChatRenderer, KimiK3ChatRenderer, }; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -73,6 +73,7 @@ impl HfChatBackend { RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), + RendererSelection::KimiK3 => Arc::new(KimiK3ChatRenderer::new(tokenizer.clone())), }; info!( diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 9c38c40be4e..0d538a24ba1 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -33,7 +33,8 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, - HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection, + HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, + RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -353,6 +354,12 @@ mod tests { .unwrap(); } + #[test] + fn validate_parser_overrides_accepts_explicit_kimi_k3() { + let selection = ParserSelection::Explicit("kimi_k3".to_string()); + validate_parser_overrides(&selection, &selection).unwrap(); + } + #[test] fn validate_parser_overrides_accepts_auto_and_none() { validate_parser_overrides(&ParserSelection::Auto, &ParserSelection::None).unwrap(); @@ -366,7 +373,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -377,6 +384,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index ca025268b04..9a1e37533b4 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -27,6 +27,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const MINIMAX_M2: &str = "minimax_m2"; pub const MINIMAX_M3: &str = "minimax_m3"; pub const NEMOTRON_V3: &str = "nemotron_v3"; @@ -68,6 +69,7 @@ impl ReasoningParserFactory { .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::MINIMAX_M2) .register_parser::(names::MINIMAX_M3) .register_parser::(names::NEMOTRON_V3) diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 78ccde8a4b3..e42b45a4862 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -33,6 +33,7 @@ pub mod names { // also routes to `Internlm2ToolParser` despite the version-agnostic name. pub const INTERNLM: &str = "internlm"; pub const KIMI_K2: &str = "kimi_k2"; + pub const KIMI_K3: &str = "kimi_k3"; pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; pub const MINIMAX_M2: &str = "minimax_m2"; @@ -78,6 +79,7 @@ impl ToolParserFactory { .register_parser::(names::HY_V3) .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) + .register_unified_dummy(names::KIMI_K3) .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) .register_parser::(names::MINIMAX_M2) diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 2536bf78c25..6246733db28 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -5,7 +5,9 @@ use std::sync::LazyLock; -pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser}; +pub use vllm_parser::unified::{ + Gemma4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, UnifiedParser, +}; use vllm_tokenizer::DynTokenizer; use crate::parser::ParserFactory; @@ -15,6 +17,7 @@ use crate::request::ChatTool; pub mod names { pub const GEMMA4: &str = "gemma4"; pub const INKLING: &str = "inkling"; + pub const KIMI_K3: &str = "kimi_k3"; } /// Constructor signature for one registered unified parser implementation. @@ -39,11 +42,14 @@ impl UnifiedParserFactory { factory.register_parser::(names::GEMMA4); factory.register_parser::(names::INKLING); + factory.register_parser::(names::KIMI_K3); factory .register_pattern("gemma-4", names::GEMMA4) .register_pattern("gemma4", names::GEMMA4) - .register_pattern("inkling", names::INKLING); + .register_pattern("inkling", names::INKLING) + .register_pattern("kimi-k3", names::KIMI_K3) + .register_pattern("kimi_k3", names::KIMI_K3); factory } @@ -121,4 +127,20 @@ mod tests { ); factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap(); } + + #[test] + fn factory_registers_kimi_k3() { + let factory = UnifiedParserFactory::new(); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|open|>", 1001) + .with_regular_token("<|close|>", 1002) + .with_regular_token("<|sep|>", 1003); + + assert!(factory.contains(names::KIMI_K3)); + assert_eq!( + factory.resolve_name_for_model("moonshotai/Kimi-K3"), + Some(names::KIMI_K3) + ); + factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap(); + } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 76796edd096..7960927219b 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -59,7 +59,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 73380cef260..9068d460ee6 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -27,7 +27,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fn deepseek_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: true, + enable_thinking: Some(true), no_generation_prompt_when_last_assistant: true, } } diff --git a/rust/src/chat/src/renderer/harmony/tests.rs b/rust/src/chat/src/renderer/harmony/tests.rs index 301d9707e71..4849c932233 100644 --- a/rust/src/chat/src/renderer/harmony/tests.rs +++ b/rust/src/chat/src/renderer/harmony/tests.rs @@ -22,7 +22,7 @@ fn fixture_request(input_name: &str) -> ChatRequest { fixture_chat_request( &fixture_path(input_name), FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, }, ) diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs index 57e08ac0158..1c4cd7236db 100644 --- a/rust/src/chat/src/renderer/inkling/tests.rs +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -105,7 +105,7 @@ fn fixture_request(name: &str) -> ChatRequest { fn inkling_fixture_options() -> FixtureRequestOptions { FixtureRequestOptions { - enable_thinking: false, + enable_thinking: Some(false), no_generation_prompt_when_last_assistant: false, } } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs new file mode 100644 index 00000000000..189d68c101c --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -0,0 +1,608 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Kimi K3 XTML prompt renderer. +//! +//! Port of Moonshot remote-code `encoding_k3.py::build_chat_segments()`. + +use std::collections::HashMap; + +use serde_json::{Map, Value, json}; +use vllm_tokenizer::Tokenizer; + +use crate::error::{Error, Result}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, +}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +pub(super) const OPEN: &str = "<|open|>"; +pub(super) const CLOSE: &str = "<|close|>"; +pub(super) const SEP: &str = "<|sep|>"; +pub(super) const END_OF_MSG: &str = "<|end_of_msg|>"; +pub(super) const IMAGE_PLACEHOLDER: &str = "<|media_pad|>"; + +const DEFAULT_THINKING_EFFORT: &str = "max"; +const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; + +/// K3 prompt encoder preserving Python's per-segment tokenization boundaries. +pub(super) struct K3TokenWriter<'a> { + tokenizer: &'a dyn Tokenizer, + token_ids: Vec, +} + +impl<'a> K3TokenWriter<'a> { + pub(super) fn new(tokenizer: &'a dyn Tokenizer) -> Self { + Self { + tokenizer, + token_ids: Vec::new(), + } + } + + /// Encode one trusted segment with normal added-token recognition. + pub(super) fn control(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode(text, false)?); + } + Ok(()) + } + + /// Encode one literal segment while bypassing every added-token matcher. + pub(super) fn ordinary(&mut self, text: &str) -> Result<()> { + if !text.is_empty() { + self.token_ids.extend(self.tokenizer.encode_ordinary(text)?); + } + Ok(()) + } + + pub(super) fn finish(self) -> Vec { + self.token_ids + } +} + +/// Render and tokenize one chat request using K3's segment-aware contract. +pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { + let thinking = thinking_enabled(request)?; + let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; + let tools = request_tools(request); + let mut out = K3TokenWriter::new(tokenizer); + + if !tools.is_empty() { + write_tool_declare(&mut out, tools, false)?; + } + + if let Some(effort) = thinking_effort { + // Preserve the checkpoint's literal guidance text: it still names + // `medium`, although the validator above no longer accepts it. + write_internal_system( + &mut out, + "thinking-effort", + &format!( + "`thinking_effort` guides on how much to think in your \ + thinking channel (not including the response channel), \ + supported values include `low`, `medium`, `high`, and `max`.\n\ + Now the system is invoked with `thinking_effort={effort}`." + ), + )?; + } + + // Track prior assistant tool-call ids for tool-result reordering / naming. + let mut tool_call_id_index: HashMap = HashMap::new(); + let mut pending_tool_run: Vec<(usize, ChatMessage)> = Vec::new(); + + let flush_tool_run = |out: &mut K3TokenWriter<'_>, + run: &mut Vec<(usize, ChatMessage)>, + id_index: &HashMap| + -> Result<()> { + if run.is_empty() { + return Ok(()); + } + + let mut resolved = Vec::with_capacity(run.len()); + let mut unresolved = false; + for (offset, message) in run.drain(..) { + let ChatMessage::ToolResponse { + content, + tool_call_id, + } = message + else { + unreachable!("pending tool run only holds tool responses"); + }; + match id_index.get(&tool_call_id) { + Some(&(position, ref name)) => { + resolved.push((position, offset, content, Some(name.clone()))); + } + None => { + unresolved = true; + resolved.push((usize::MAX, offset, content, None)); + } + } + } + + if unresolved { + // Preserve original order when the run cannot be fully matched. + resolved.sort_by_key(|item| item.1); + } else { + resolved.sort_by_key(|item| (item.0, item.1)); + } + + for (xtml_index, (_, _, content, name)) in resolved.into_iter().enumerate() { + let tool_name = name.as_deref().ok_or_else(|| { + Error::ChatTemplate( + "Kimi K3 tool messages need a resolvable tool name: \ + carry a matching tool_call_id against a preceding \ + assistant tool_call" + .to_string(), + ) + })?; + write_tool_message(out, tool_name, xtml_index + 1, &content)?; + } + Ok(()) + }; + + for (message_index, message) in request.messages.iter().enumerate() { + match message { + ChatMessage::ToolResponse { .. } => { + pending_tool_run.push((message_index, message.clone())); + continue; + } + _ => { + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + } + } + + match message { + // Python: role=system with a `tools` field renders a dynamic + // tool-declare (`## New Tools Available`). Map that to Developer + // messages that carry tools (OpenAI "developer" / system-tools). + ChatMessage::Developer { + content, + tools: Some(local_tools), + } if !local_tools.is_empty() => { + write_tool_declare(&mut out, local_tools, true)?; + if !content_is_empty(content) { + write_role_message(&mut out, "system", None, content)?; + } + } + ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { + write_role_message(&mut out, "system", None, content)?; + } + ChatMessage::User { content } => { + write_role_message(&mut out, "user", None, content)?; + } + ChatMessage::Assistant { content } => { + tool_call_id_index.clear(); + let mut call_position = 0usize; + for block in content { + if let AssistantContentBlock::ToolCall(call) = block { + call_position += 1; + if !call.id.is_empty() { + tool_call_id_index + .entry(call.id.clone()) + .or_insert((call_position, call.name.clone())); + } + } + } + + write_assistant_message(&mut out, content, thinking)?; + } + ChatMessage::ToolResponse { .. } => unreachable!("handled above"), + } + } + flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; + + match &request.tool_choice { + ChatToolChoice::Required => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=required`.\n\ + You MUST call tools in the next message.", + )?; + } + // Emit only when tools are present: Rust defaults tool_choice to None + // for tool-free requests, which must not inject a tool-choice message. + ChatToolChoice::None if !request.tools.is_empty() => { + write_internal_system( + &mut out, + "tool-choice", + "The system is invoked with `tool_choice=none`.\n\ + You MUST NOT call any tools in the next message.", + )?; + } + ChatToolChoice::None | ChatToolChoice::Auto | ChatToolChoice::Function { .. } => {} + } + + write_response_format(&mut out, request)?; + + if request.chat_options.add_generation_prompt() { + write_open_tag(&mut out, "message", &[("role", "assistant")])?; + write_open_tag(&mut out, if thinking { "think" } else { "response" }, &[])?; + } + + Ok(out.finish()) +} + +fn request_tools(request: &ChatRequest) -> &[ChatTool] { + // Declare tools whenever the request carries them. K3 tool-declare is + // independent of tool_choice; tool_choice only injects control messages. + request.tools.as_slice() +} + +fn thinking_enabled(request: &ChatRequest) -> Result { + if let Some(thinking) = request.parse_template_bool("thinking")? { + return Ok(thinking); + } + if let Some(enable_thinking) = request.parse_template_bool("enable_thinking")? { + return Ok(enable_thinking); + } + Ok(request + .chat_options + .reasoning_effort + .map(|effort| effort != crate::request::ReasoningEffort::None) + .unwrap_or(true)) +} + +fn thinking_effort(request: &ChatRequest) -> Result { + let effort = if let Some(value) = request.chat_options.template_kwargs.get("thinking_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `thinking_effort` must be a string, got {value}" + )) + })? + } else if let Some(effort) = request.chat_options.reasoning_effort { + effort.as_str() + } else if let Some(value) = request.chat_options.template_kwargs.get("reasoning_effort") { + value.as_str().ok_or_else(|| { + Error::ChatTemplate(format!( + "template kwarg `reasoning_effort` must be a string, got {value}" + )) + })? + } else { + DEFAULT_THINKING_EFFORT + }; + + if !VALID_THINKING_EFFORTS.contains(&effort) { + return Err(Error::ChatTemplate(format!( + "unsupported thinking_effort={effort:?}; supported values are `low`, `high`, and `max`" + ))); + } + Ok(effort.to_string()) +} + +fn content_is_empty(content: &ChatContent) -> bool { + match content { + ChatContent::Text(text) => text.is_empty(), + ChatContent::Parts(parts) => parts.iter().all(|part| match part { + ChatContentPart::Text { text } => text.is_empty(), + ChatContentPart::ImageUrl { .. } + | ChatContentPart::VideoUrl { .. } + | ChatContentPart::InputAudio { .. } + | ChatContentPart::AudioUrl { .. } => false, + }), + } +} + +fn write_tool_declare( + out: &mut K3TokenWriter<'_>, + tools: &[ChatTool], + dynamic: bool, +) -> Result<()> { + let mut specs = Vec::with_capacity(tools.len()); + for tool in tools { + let mut function = Map::new(); + function.insert( + "description".to_string(), + Value::String(tool.description.clone().unwrap_or_default()), + ); + function.insert("name".to_string(), Value::String(tool.name.clone())); + function.insert("parameters".to_string(), sort_json(&tool.parameters)); + specs.push(json!({ + "function": Value::Object(function), + "type": "function", + })); + } + let payload = compact_json(&sort_json(&Value::Array(specs)))?; + + let body = if dynamic { + format!( + "## New Tools Available\n\ + The system dynamically extends the toolset via lazy-loading.\n\ + You have access to all existing and extended tools.\n\ + Here are the specs for the extended tools.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + } else { + format!( + "# Tools\n\ + Here are the available tools, described in JSONSchema.\n\n\ + ```json\n\ + {payload}\n\ + ```" + ) + }; + + write_internal_system(out, "tool-declare", &body) +} + +fn write_internal_system( + out: &mut K3TokenWriter<'_>, + message_type: &str, + body: &str, +) -> Result<()> { + write_open_tag( + out, + "message", + &[("role", "system"), ("type", message_type)], + )?; + out.ordinary(body.trim())?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_role_message( + out: &mut K3TokenWriter<'_>, + role: &str, + name: Option<&str>, + content: &ChatContent, +) -> Result<()> { + let mut attrs = vec![("role", role.to_string())]; + if let Some(name) = name.filter(|name| !name.is_empty()) { + attrs.push(("name", name.to_string())); + } + let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); + write_open_tag(out, "message", &attr_refs)?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_tool_message( + out: &mut K3TokenWriter<'_>, + tool_name: &str, + index: usize, + content: &ChatContent, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "message", + &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], + )?; + write_content(out, content)?; + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_message( + out: &mut K3TokenWriter<'_>, + content: &[AssistantContentBlock], + thinking: bool, +) -> Result<()> { + write_open_tag(out, "message", &[("role", "assistant")])?; + + let mut reasoning = String::new(); + let mut response = String::new(); + let mut tool_calls = Vec::new(); + for block in content { + match block { + AssistantContentBlock::Reasoning { text } => reasoning.push_str(text), + AssistantContentBlock::Text { text } => response.push_str(text), + AssistantContentBlock::ToolCall(call) => tool_calls.push(call), + } + } + + // The think channel is structural: in thinking mode every assistant + // message carries open/close tags even when there is no reasoning content. + // In non-thinking mode the channel is dropped entirely. + if thinking { + write_open_tag(out, "think", &[])?; + if !reasoning.trim().is_empty() { + out.ordinary(&reasoning)?; + } + write_close_tag(out, "think")?; + } + + write_open_tag(out, "response", &[])?; + out.ordinary(&response)?; + write_close_tag(out, "response")?; + + if !tool_calls.is_empty() { + write_open_tag(out, "tools", &[])?; + for (index, tool_call) in tool_calls.into_iter().enumerate() { + write_assistant_tool_call(out, tool_call, index + 1)?; + } + write_close_tag(out, "tools")?; + } + + write_close_tag(out, "message")?; + out.control(END_OF_MSG) +} + +fn write_assistant_tool_call( + out: &mut K3TokenWriter<'_>, + tool_call: &AssistantToolCall, + index: usize, +) -> Result<()> { + let index_str = index.to_string(); + write_open_tag( + out, + "call", + &[ + ("tool", tool_call.name.as_str()), + ("index", index_str.as_str()), + ], + )?; + + let (args, json_block) = normalize_tool_arguments(&tool_call.arguments)?; + if let Some(raw) = json_block { + write_open_tag(out, "json", &[("type", "object")])?; + out.ordinary(&raw)?; + write_close_tag(out, "json")?; + } else { + for (key, value) in args { + let typ = xtml_type(&value); + write_open_tag(out, "argument", &[("key", key.as_str()), ("type", typ)])?; + out.ordinary(&xtml_value(&value))?; + write_close_tag(out, "argument")?; + } + } + + write_close_tag(out, "call") +} + +fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { + match content { + ChatContent::Text(text) => write_text_with_images(out, text), + ChatContent::Parts(parts) => { + for part in parts { + match part { + ChatContentPart::Text { text } => write_text_with_images(out, text)?, + ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::VideoUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("video_url")); + } + ChatContentPart::InputAudio { .. } => { + return Err(Error::UnsupportedMultimodalContent("input_audio")); + } + ChatContentPart::AudioUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("audio_url")); + } + } + } + Ok(()) + } + } +} + +fn write_text_with_images(out: &mut K3TokenWriter<'_>, text: &str) -> Result<()> { + // Placeholder expansion is left as the literal K3 image token; multimodal + // preprocessing can replace it once image prompts are known. + out.ordinary(text) +} + +fn write_response_format(out: &mut K3TokenWriter<'_>, request: &ChatRequest) -> Result<()> { + let Some(rf) = request.chat_options.response_format.as_ref() else { + return Ok(()); + }; + + let rf_type = rf.get("type").and_then(Value::as_str).or_else(|| rf.as_str()).unwrap_or(""); + + match rf_type { + "json_object" => { + write_internal_system( + out, + "response-format", + "The system is invoked with `response_format=json_object`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.", + )?; + } + "json_schema" => { + let schema = extract_response_schema(rf); + let schema_json = compact_json(&sort_json(&schema.unwrap_or(Value::Null)))?; + write_internal_system( + out, + "response-format", + &format!( + "The system is invoked with `response_format=json_schema`.\n\ + Your response must be raw JSON data without markdown code \ + blocks (```json) or any additional formatting.\n\ + The JSON data must match the following schema:\n\ + ```json\n\ + {schema_json}\n\ + ```" + ), + )?; + } + _ => {} + } + Ok(()) +} + +fn extract_response_schema(response_format: &Value) -> Option { + let json_schema = response_format.get("json_schema")?; + if let Some(schema) = json_schema.get("schema") { + return Some(schema.clone()); + } + if let Some(schema) = json_schema.get("json_schema") { + return Some(schema.clone()); + } + Some(json_schema.clone()) +} + +fn normalize_tool_arguments(arguments: &str) -> Result<(Map, Option)> { + let trimmed = arguments.trim(); + if trimmed.is_empty() { + return Ok((Map::new(), None)); + } + match serde_json::from_str::(trimmed) { + Ok(Value::Object(map)) => Ok((map, None)), + Ok(_) => Err(Error::ChatTemplate( + "Kimi K3 tool call arguments must be a JSON object".to_string(), + )), + Err(_) => Ok((Map::new(), Some(arguments.to_string()))), + } +} + +fn xtml_type(value: &Value) -> &'static str { + match value { + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Object(_) => "object", + Value::Array(_) => "array", + } +} + +fn xtml_value(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + other => compact_json(other).unwrap_or_else(|_| other.to_string()), + } +} + +fn write_open_tag(out: &mut K3TokenWriter<'_>, tag: &str, attrs: &[(&str, &str)]) -> Result<()> { + out.control(OPEN)?; + out.ordinary(tag)?; + for (key, value) in attrs { + out.ordinary(&format!(" {key}"))?; + out.ordinary("=\"")?; + out.ordinary(&escape_attr_value(value))?; + out.ordinary("\"")?; + } + out.control(SEP) +} + +fn write_close_tag(out: &mut K3TokenWriter<'_>, tag: &str) -> Result<()> { + out.control(CLOSE)?; + out.ordinary(tag)?; + out.control(SEP) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +fn compact_json(value: &Value) -> Result { + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.to_string())) +} + +fn sort_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()), + Value::Object(map) => { + let mut sorted = Map::new(); + let mut keys = map.keys().collect::>(); + keys.sort(); + for key in keys { + sorted.insert(key.clone(), sort_json(&map[key])); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json new file mode 100644 index 00000000000..569201f3b97 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_input.json @@ -0,0 +1,15 @@ +{ + "messages": [ + { + "role": "user", + "content": "json pls" + } + ], + "add_generation_prompt": true, + "response_format": { + "type": "json_object" + }, + "template_kwargs": { + "thinking": false + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt new file mode 100644 index 00000000000..57ce4a68761 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/controls_thinking_off_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="user"<|sep|>json pls<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="response-format"<|sep|>The system is invoked with `response_format=json_object`. +Your response must be raw JSON data without markdown code blocks (```json) or any additional formatting.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>response<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json new file mode 100644 index 00000000000..61ecebc268b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_input.json @@ -0,0 +1,62 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "developer", + "tools": [ + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ] + }, + { + "role": "user", + "content": "use calc" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt new file mode 100644 index 00000000000..edec2bebf2a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_system_tool_declare_output.txt @@ -0,0 +1,14 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>hi<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available +The system dynamically extends the toolset via lazy-loading. +You have access to all existing and extended tools. +Here are the specs for the extended tools. + +```json +[{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>use calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json new file mode 100644 index 00000000000..b7cf40c9b8b --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_input.json @@ -0,0 +1,31 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "see this" + }, + { + "type": "image_url", + "image_url": "https://example.com/a.png" + } + ] + }, + { + "role": "assistant", + "content": "old answer", + "reasoning_content": "old reasoning" + }, + { + "role": "user", + "content": "continue" + } + ], + "add_generation_prompt": true, + "template_kwargs": { + "thinking": true, + "thinking_effort": "high" + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt new file mode 100644 index 00000000000..03582f1180a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/history_preserve_and_image_output.txt @@ -0,0 +1,2 @@ +<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=high`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>see this<|media_pad|><|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>old reasoning<|close|>think<|sep|><|open|>response<|sep|>old answer<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>continue<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json new file mode 100644 index 00000000000..3a64f8dd0bb --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_input.json @@ -0,0 +1,87 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "weather", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string" + }, + "days": { + "type": "number" + } + }, + "required": [ + "city" + ] + } + } + }, + { + "type": "function", + "function": { + "name": "calc", + "description": "calculator", + "parameters": { + "type": "object", + "properties": { + "x": { + "type": "number" + } + }, + "required": [ + "x" + ] + } + } + } + ], + "messages": [ + { + "role": "user", + "content": "Hangzhou weather and calc" + }, + { + "role": "assistant", + "content": "I'll check.", + "reasoning_content": "Need tools.", + "tool_calls": [ + { + "id": "get_weather:0", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"Hangzhou\",\"days\":1}" + } + }, + { + "id": "calc:1", + "type": "function", + "function": { + "name": "calc", + "arguments": "{\"x\":2}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "calc:1", + "content": "2" + }, + { + "role": "tool", + "tool_call_id": "get_weather:0", + "content": "{\"city\":\"Hangzhou\",\"condition\":\"rain\"}" + } + ], + "add_generation_prompt": true, + "tool_choice": "required", + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt new file mode 100644 index 00000000000..1c24810072a --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/tools_history_and_required_output.txt @@ -0,0 +1,8 @@ +<|open|>message role="system" type="tool-declare"<|sep|># Tools +Here are the available tools, described in JSONSchema. + +```json +[{"function":{"description":"weather","name":"get_weather","parameters":{"properties":{"city":{"type":"string"},"days":{"type":"number"}},"required":["city"],"type":"object"}},"type":"function"},{"function":{"description":"calculator","name":"calc","parameters":{"properties":{"x":{"type":"number"}},"required":["x"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Hangzhou weather and calc<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>Need tools.<|close|>think<|sep|><|open|>response<|sep|>I'll check.<|close|>response<|sep|><|open|>tools<|sep|><|open|>call tool="get_weather" index="1"<|sep|><|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|><|open|>argument key="days" type="number"<|sep|>1<|close|>argument<|sep|><|close|>call<|sep|><|open|>call tool="calc" index="2"<|sep|><|open|>argument key="x" type="number"<|sep|>2<|close|>argument<|sep|><|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="get_weather" index="1"<|sep|>{"city":"Hangzhou","condition":"rain"}<|close|>message<|sep|><|end_of_msg|><|open|>message role="tool" tool="calc" index="2"<|sep|>2<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=required`. +You MUST call tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs new file mode 100644 index 00000000000..09bf356f5d6 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Native Kimi K3 XTML chat renderer. + +mod encoding; +#[cfg(test)] +mod tests; + +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; + +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::Result; +use crate::request::ChatRequest; + +/// Dedicated Kimi K3 XTML renderer. +#[derive(Clone)] +pub struct KimiK3ChatRenderer { + tokenizer: DynTokenizer, +} + +impl KimiK3ChatRenderer { + /// Create a Kimi K3 renderer. + pub fn new(tokenizer: DynTokenizer) -> Self { + Self { tokenizer } + } +} + +impl ChatRenderer for KimiK3ChatRenderer { + fn render(&self, request: &ChatRequest) -> Result { + request.validate()?; + + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + effective_template_kwargs: request_template_kwargs(request), + }) + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs new file mode 100644 index 00000000000..7bab90ddcf7 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Golden fixtures generated from HF remote-code `encoding_k3.py`. + +use std::path::PathBuf; +use std::sync::Arc; + +use expect_test::{expect, expect_file}; +use serde_json::json; +use vllm_text::Prompt; +use vllm_text::tokenizer::DynTokenizer; +use vllm_tokenizer::Tokenizer; +use vllm_tokenizer::test_utils::TestTokenizer; + +use super::KimiK3ChatRenderer; +use crate::AssistantContentBlock; +use crate::ChatRenderer; +use crate::renderer::kimi_k3::encoding::{CLOSE, END_OF_MSG, IMAGE_PLACEHOLDER, OPEN, SEP}; +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; +use crate::request::{ChatContentPart, ChatMessage, GenerationPromptMode, ReasoningEffort}; + +const OPEN_ID: u32 = 256; +const CLOSE_ID: u32 = 257; +const SEP_ID: u32 = 258; +const END_OF_MSG_ID: u32 = 259; +const MEDIA_ID: u32 = 260; + +fn test_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token(CLOSE, CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + .with_special_token(IMAGE_PLACEHOLDER, MEDIA_ID) +} + +fn render_token_ids(request: &crate::request::ChatRequest, tokenizer: DynTokenizer) -> Vec { + let prompt = KimiK3ChatRenderer::new(tokenizer).render(request).unwrap().prompt; + let Prompt::TokenIds(token_ids) = prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + token_ids +} + +fn render_request(request: &crate::request::ChatRequest) -> String { + let tokenizer: DynTokenizer = Arc::new(test_tokenizer()); + let token_ids = render_token_ids(request, tokenizer.clone()); + tokenizer.decode(&token_ids, false).unwrap() +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/kimi_k3/fixtures") + .join(name) +} + +fn kimi_k3_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + // Fixture JSON owns thinking via `template_kwargs`. + enable_thinking: None, + no_generation_prompt_when_last_assistant: false, + } +} + +fn assert_golden(name: &str) { + let input_name = format!("{name}_input.json"); + let request = fixture_chat_request(&fixture_path(&input_name), kimi_k3_fixture_options()); + let rendered = render_request(&request); + expect_file![format!("fixtures/{name}_output.txt")].assert_eq(&rendered); +} + +#[test] +fn golden_history_preserve_and_image() { + assert_golden("history_preserve_and_image"); +} + +#[test] +fn golden_tools_history_and_required() { + assert_golden("tools_history_and_required"); +} + +#[test] +fn golden_controls_thinking_off() { + assert_golden("controls_thinking_off"); +} + +#[test] +fn golden_dynamic_system_tool_declare() { + assert_golden("dynamic_system_tool_declare"); +} + +#[test] +fn token_writer_protects_literal_control_and_media_markers() { + let tokenizer = Arc::new(test_tokenizer()); + let user_text = format!("literal {OPEN} and {}", super::encoding::IMAGE_PLACEHOLDER); + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ChatMessage::user(vec![ + ChatContentPart::text(user_text), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let token_ids = render_token_ids(&request, tokenizer.clone()); + + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 1 + ); + assert_eq!( + token_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 1 + ); + + let flattened = tokenizer.decode(&token_ids, false).unwrap(); + let flattened_ids = tokenizer.encode(&flattened, false).unwrap(); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == OPEN_ID).count(), + 2 + ); + assert_eq!( + flattened_ids.iter().filter(|&&token_id| token_id == MEDIA_ID).count(), + 2 + ); +} + +#[test] +fn thinking_history_renders_empty_think_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>think<|sep|><|close|>think<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn non_thinking_history_omits_reasoning_channel() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::user("question"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { + text: "hidden reasoning".to_string(), + }, + AssistantContentBlock::Text { + text: "answer".to_string(), + }, + ]), + ChatMessage::user("follow-up"), + ]; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), json!(false)); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(!rendered.contains("hidden reasoning")); + assert!(!rendered.contains("<|open|>think<|sep|>")); + assert!(rendered.contains( + "<|open|>message role=\"assistant\"<|sep|>\ + <|open|>response<|sep|>answer<|close|>response<|sep|>" + )); +} + +#[test] +fn defaults_thinking_effort_to_max() { + let rendered = render_request(&crate::request::ChatRequest::for_test()); + + expect![[r#"<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>test<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|>"#]] + .assert_eq(&rendered); +} + +#[test] +fn translates_standard_thinking_kwargs() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), json!(true)); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!("high")); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=high")); + assert!(rendered.ends_with("<|open|>think<|sep|>")); +} + +#[test] +fn native_k3_kwargs_take_precedence() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("thinking".to_string(), json!(true)), + ("enable_thinking".to_string(), json!(false)), + ("thinking_effort".to_string(), json!("low")), + ("reasoning_effort".to_string(), json!("high")), + ]); + + let rendered = render_request(&request); + + assert!(rendered.contains("thinking_effort=low")); + assert!(!rendered.contains("thinking_effort=high")); +} + +#[test] +fn standard_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.template_kwargs.extend([ + ("enable_thinking".to_string(), json!(false)), + ("reasoning_effort".to_string(), json!("none")), + ]); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn typed_none_disables_thinking() { + let mut request = crate::request::ChatRequest::for_test(); + request.chat_options.reasoning_effort = Some(ReasoningEffort::None); + + let rendered = render_request(&request); + + assert!(!rendered.contains("type=\"thinking-effort\"")); + assert!(rendered.ends_with("<|open|>response<|sep|>")); +} + +#[test] +fn rejects_removed_medium_thinking_effort() { + let mut request = crate::request::ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("thinking_effort".to_string(), json!("medium")); + + let error = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())) + .render(&request) + .unwrap_err(); + + expect![[r#" + ChatTemplate( + "unsupported thinking_effort=\"medium\"; supported values are `low`, `high`, and `max`", + ) + "#]] + .assert_debug_eq(&error); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index 4d9c1581a01..a83953f3318 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -15,6 +15,7 @@ pub mod deepseek_v4; pub mod harmony; pub mod hf; mod inkling; +mod kimi_k3; mod selection; #[cfg(test)] mod test_utils; @@ -23,6 +24,7 @@ pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; pub use harmony::HarmonyChatRenderer; pub use inkling::InklingChatRenderer; +pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 711c0d87ddf..5b6642c4c03 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -26,6 +26,8 @@ pub enum RendererSelection { Harmony, /// Force the Inkling native token renderer. Inkling, + /// Force the Kimi K3 XTML renderer. + KimiK3, } impl RendererSelection { @@ -37,6 +39,7 @@ impl RendererSelection { pub const HF_LITERAL: &str = "hf"; pub const INKLING_LITERAL: &str = "inkling"; pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model"; + pub const KIMI_K3_LITERAL: &str = "kimi_k3"; /// Resolve the renderer selection using the given model type string, if /// it's `Auto`. @@ -47,6 +50,7 @@ impl RendererSelection { Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, Self::GPT_OSS_MODEL_TYPE => Self::Harmony, Self::INKLING_MODEL_TYPE => Self::Inkling, + Self::KIMI_K3_LITERAL => Self::KimiK3, _ => Self::Hf, }, selection => selection, @@ -70,6 +74,8 @@ impl FromStr for RendererSelection { Ok(Self::Harmony) } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) { Ok(Self::Inkling) + } else if value.eq_ignore_ascii_case(Self::KIMI_K3_LITERAL) { + Ok(Self::KimiK3) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -88,6 +94,7 @@ impl fmt::Display for RendererSelection { Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), Self::Harmony => f.write_str(Self::HARMONY_LITERAL), Self::Inkling => f.write_str(Self::INKLING_LITERAL), + Self::KimiK3 => f.write_str(Self::KIMI_K3_LITERAL), } } } @@ -114,7 +121,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 08eeca8176d..0a4edc81379 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::HashMap; use std::fs; use std::path::Path; @@ -16,8 +17,9 @@ use crate::request::{ /// Options for constructing a [`ChatRequest`] from a fixture file. #[derive(Debug, Clone, Copy)] pub(crate) struct FixtureRequestOptions { - /// Whether to set the template kwarg `[enable_]thinking=true`. - pub enable_thinking: bool, + /// Optional thinking toggle applied only when the fixture does not already + /// set `thinking` / `enable_thinking` in [`FixtureRequest::template_kwargs`]. + pub enable_thinking: Option, /// Whether fixtures ending in an assistant message should omit the /// trailing generation prompt. pub no_generation_prompt_when_last_assistant: bool, @@ -46,6 +48,15 @@ pub(crate) struct FixtureRequest { messages: Vec, add_generation_prompt: Option, reasoning_effort: Option, + /// Standard response format passed to model-specific renderers. + #[serde(default)] + response_format: Option, + /// Extra chat-template kwargs (thinking, preserve_thinking, …). + #[serde(default)] + template_kwargs: HashMap, + /// When omitted, defaults to `auto` if tools are present, otherwise `none`. + #[serde(default)] + tool_choice: Option, } impl FixtureFile { @@ -57,6 +68,9 @@ impl FixtureFile { messages, add_generation_prompt: None, reasoning_effort: None, + response_format: None, + template_kwargs: HashMap::new(), + tool_choice: None, }, } } @@ -69,6 +83,7 @@ pub(crate) enum FixtureMessage { content: FixtureContent, }, Developer { + #[serde(default)] content: FixtureContent, #[serde(default)] tools: Vec, @@ -98,6 +113,12 @@ pub(crate) enum FixtureContent { Parts(Vec), } +impl Default for FixtureContent { + fn default() -> Self { + Self::Text(String::new()) + } +} + #[derive(Debug, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub(crate) enum FixtureContentPart { @@ -136,6 +157,13 @@ struct FixtureToolCallFunction { impl FixtureRequest { fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { + let tools = to_chat_tools(&self.tools); + let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + }); + let mut request = ChatRequest { request_id: "renderer-fixture".to_string(), messages: self @@ -144,12 +172,8 @@ impl FixtureRequest { .enumerate() .map(|(index, message)| fixture_message_to_chat_message(index, message)) .collect(), - tools: to_chat_tools(&self.tools), - tool_choice: if self.tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, + tools, + tool_choice, ..ChatRequest::for_test() }; @@ -162,9 +186,15 @@ impl FixtureRequest { request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; } request.chat_options.reasoning_effort = self.reasoning_effort; - if options.enable_thinking { - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), Value::Bool(true)); + request.chat_options.response_format = self.response_format; + request.chat_options.template_kwargs.extend(self.template_kwargs); + + // Options supply a default thinking toggle only when the fixture did not. + if let Some(thinking) = options.enable_thinking { + let kwargs = &mut request.chat_options.template_kwargs; + if !kwargs.contains_key("thinking") && !kwargs.contains_key("enable_thinking") { + kwargs.insert("thinking".to_string(), Value::Bool(thinking)); + kwargs.insert("enable_thinking".to_string(), Value::Bool(thinking)); } } diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index eb72557c7d2..56a764f8987 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -397,6 +397,10 @@ pub struct ChatOptions { /// Effort level exposed to chat templates for reasoning models. pub reasoning_effort: Option, + /// Standard response format available to model-specific renderers. + #[serde(default)] + pub response_format: Option, + /// Additional keyword arguments exposed to the chat template. pub template_kwargs: HashMap, } @@ -407,6 +411,7 @@ impl Default for ChatOptions { generation_prompt_mode: GenerationPromptMode::StartNewAssistant, chat_template: None, reasoning_effort: None, + response_format: None, template_kwargs: HashMap::new(), } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index ff15aa929ef..9cd670d29d0 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -212,6 +212,23 @@ impl RoundtripCase { } } + /// Kimi K3 XTML tool/reasoning channels (native renderer + unified parser). + /// + /// Needs HF tokenizer files under `HF_HOME` (`tiktoken.model` + + /// `tokenizer_config.json`). Weights are not required for this text-level + /// roundtrip. + fn kimi_k3() -> Self { + Self { + model_id: "moonshotai/Kimi-K3", + assistant_stop_suffix: "<|end_of_msg|>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// SeedOSS with `` / `` reasoning tags. fn seed_oss() -> Self { Self { @@ -255,7 +272,7 @@ impl RoundtripCase { fn gpt_oss() -> Self { Self { model_id: "openai/gpt-oss-20b", - assistant_stop_suffix: "", // not applicable for token-id cases + assistant_stop_suffix: "", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -268,7 +285,7 @@ impl RoundtripCase { fn inkling() -> Self { Self { model_id: "thinkingmachines/Inkling", - assistant_stop_suffix: "", + assistant_stop_suffix: "<|content_model_end_sampling|>", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, thinking_behavior: ThinkingBehavior::Always { value: true }, @@ -312,6 +329,8 @@ roundtrip_tests! { nemotron_v3 => [reasoning_and_content], gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history + // K3 drops plain-assistant reasoning in history; tool-call turns keep it. + kimi_k3 => [tool_call_mix], gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call inkling => [reasoning_and_content, tool_call_mix], } @@ -670,14 +689,23 @@ fn decoded_completion_stream( .collect() } Prompt::TokenIds(token_ids) => { - ensure!( - assistant_stop_suffix.is_empty(), - "token-id roundtrip cases do not support text stop suffixes" - ); + let body = if assistant_stop_suffix.is_empty() { + token_ids.as_slice() + } else { + let stop_token_ids = tokenizer + .encode(assistant_stop_suffix, false) + .context("failed to encode token-id completion stop suffix")?; + token_ids.strip_suffix(stop_token_ids.as_slice()).with_context(|| { + format!( + "token-id completion did not end with {:?}: {:?}", + assistant_stop_suffix, token_ids + ) + })? + }; incremental_decode_chunks( tokenizer, &prompt_token_ids, - token_ids, + body, TOKEN_COMPLETION_CHUNK_TOKENS, )? } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index b0926fb24ac..bfece8b6092 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -666,7 +666,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling, kimi_k3) For more information, try '--help'. "#]] diff --git a/rust/src/parser/src/unified/kimi_k3.rs b/rust/src/parser/src/unified/kimi_k3.rs new file mode 100644 index 00000000000..7e716225885 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3.rs @@ -0,0 +1,1149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for the Kimi K3 (XTML) chat format. +//! +//! Original Python implementations: +//! - `vllm/reasoning/kimi_k3_reasoning_parser.py` +//! - `vllm/tool_parsers/kimi_k3_tool_parser.py` +//! +//! K3 wraps one assistant message into XTML channels built from the dedicated +//! special tokens `<|open|>`, `<|close|>`, and `<|sep|>`: +//! +//! ```text +//! <|open|>think<|sep|>reasoning<|close|>think<|sep|> +//! <|open|>response<|sep|>visible answer<|close|>response<|sep|> +//! <|open|>tools<|sep|> +//! <|open|>call tool="get_weather" index="1"<|sep|> +//! <|open|>argument key="city" type="string"<|sep|>Hangzhou<|close|>argument<|sep|> +//! <|close|>call<|sep|> +//! <|close|>tools<|sep|> +//! <|close|>message<|sep|> +//! ``` +//! +//! In chat serving the generation prompt ends with `<|open|>think<|sep|>` +//! (thinking) or `<|open|>response<|sep|>` (instruct), so the model output +//! starts *inside* that channel without re-emitting the open tag; +//! [`UnifiedParser::initialize`] detects this from the prompt token IDs. +//! +//! Argument decoding mirrors the renderer's type tagging (inverse encoding): +//! `type="string"` values pass the raw text through, other types are +//! JSON-decoded, and a raw `json` block is passed through unmodified. Attribute +//! values reverse the renderer escaping (`"` before `&`). +//! +//! Known limitation (shared with the Python parser): string argument and +//! response bodies are emitted raw, so a value that literally contains +//! `<|close|>argument<|sep|>` or `<|close|>response<|sep|>` is +//! indistinguishable from a real closing marker. + +mod structural_tag; + +pub use structural_tag::KimiK3StructuralTagBuilder; + +use serde_json::{Map, Value}; +use vllm_tokenizer::DynTokenizer; +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1}; +use winnow::combinator::{alt, delimited, eof, preceded, repeat, seq, terminated}; +use winnow::error::{ContextError, ErrMode, ModalResult, StrContext}; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::{literal, rest, take_till, take_until, take_while}; + +use self::structural_tag::KIMI_K3_STRUCTURAL_TAG_BUILDER; +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::{StructuralTagBuilder, Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; + +const OPEN: &str = "<|open|>"; +const SEP: &str = "<|sep|>"; +const END_OF_MSG: &str = "<|end_of_msg|>"; + +const THINK_OPEN: &str = "<|open|>think<|sep|>"; +const THINK_CLOSE: &str = "<|close|>think<|sep|>"; +const RESPONSE_OPEN: &str = "<|open|>response<|sep|>"; +const RESPONSE_CLOSE: &str = "<|close|>response<|sep|>"; +const TOOLS_OPEN: &str = "<|open|>tools<|sep|>"; +const TOOLS_CLOSE: &str = "<|close|>tools<|sep|>"; +const MESSAGE_CLOSE: &str = "<|close|>message<|sep|>"; +const CALL_OPEN: &str = "<|open|>call"; +const CALL_CLOSE: &str = "<|close|>call<|sep|>"; +const ARG_OPEN: &str = "<|open|>argument"; +const ARG_CLOSE: &str = "<|close|>argument<|sep|>"; +const JSON_OPEN: &str = "<|open|>json"; +const JSON_CLOSE: &str = "<|close|>json<|sep|>"; + +const IDLE_MARKERS: &[&str] = &[ + THINK_OPEN, + RESPONSE_OPEN, + TOOLS_OPEN, + MESSAGE_CLOSE, + END_OF_MSG, +]; +const REASONING_MARKERS: &[&str] = &[THINK_CLOSE, END_OF_MSG]; +const RESPONSE_MARKERS: &[&str] = &[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const EPILOGUE_MARKERS: &[&str] = &[TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]; +const TOOLS_MARKERS: &[&str] = &[CALL_OPEN, TOOLS_CLOSE, MESSAGE_CLOSE, END_OF_MSG]; + +/// Channel tags are a couple of text tokens; longer `<|open|>…<|sep|>` spans in +/// the prompt tail (attribute-bearing message opens, message bodies) never name +/// a generation channel. +const MAX_PREFILL_TAG_TOKENS: usize = 8; + +type KimiK3Input<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +enum KimiK3Event { + Text { + len: usize, + }, + Reasoning { + len: usize, + }, + /// Structural noise consumed without emitting anything. + Skip, + ThinkOpen, + ThinkClose, + ResponseOpen, + ResponseClose, + ToolsOpen, + ToolsClose, + /// The assistant message closed; everything after it is ignored. + MessageEnd, + CallOpen { + name: String, + index: Option, + }, + CallComplete { + arguments: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +enum KimiK3Mode { + /// Before any channel opens: channel opens are expected, but raw text + /// falls through as visible text so marker-free output still streams. + #[default] + Idle, + /// Inside the `think` channel. + Reasoning, + /// Inside the `response` channel. + Response, + /// After the `response` (or `tools`) channel closed: only a `tools` + /// channel or the message close may follow, and noise is never content. + Epilogue, + /// Inside the `tools` channel, between `call` blocks. + Tools, + /// Inside one `call` block, buffering its body until the close marker. + Call { + name: String, + index: Option, + scan: MarkerScanState, + }, + /// After the message closed: ignore the rest (EOS leakage guard). + Done, +} + +/// Unified parser for Kimi K3 XTML think / response / tools channels. +pub struct KimiK3UnifiedParser { + buffer: String, + mode: KimiK3Mode, + /// Parser-provided tool-call IDs (`{tool}:{zero_based_index}`) by tool + /// index; its length is also the count of emitted calls. + call_ids: Vec, + tokenizer: DynTokenizer, + open_token_id: u32, + sep_token_id: u32, +} + +impl KimiK3UnifiedParser { + /// Create a Kimi K3 parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let open_token_id = token_id(tokenizer.as_ref(), OPEN)?; + let sep_token_id = token_id(tokenizer.as_ref(), SEP)?; + + Ok(Self { + buffer: String::new(), + mode: KimiK3Mode::default(), + call_ids: Vec::new(), + tokenizer, + open_token_id, + sep_token_id, + }) + } + + /// Detect the prefilled generation channel from the prompt tail. + /// + /// `add_generation_prompt` ends the prompt with `<|open|>think<|sep|>` or + /// `<|open|>response<|sep|>`, so generation starts inside that channel and + /// never re-emits the open tag. Locate the last `<|open|>…<|sep|>` pair and + /// decode the short tag between the two structural tokens. + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = KimiK3Mode::Idle; + + let Some(sep_pos) = prompt_token_ids.iter().rposition(|&id| id == self.sep_token_id) else { + return; + }; + let Some(open_pos) = + prompt_token_ids[..sep_pos].iter().rposition(|&id| id == self.open_token_id) + else { + return; + }; + let tag_ids = &prompt_token_ids[open_pos + 1..sep_pos]; + if tag_ids.is_empty() || tag_ids.len() > MAX_PREFILL_TAG_TOKENS { + return; + } + let Ok(tag) = self.tokenizer.decode(tag_ids, /* skip_special_tokens */ false) else { + return; + }; + self.mode = match tag.trim() { + "think" => KimiK3Mode::Reasoning, + "response" => KimiK3Mode::Response, + _ => KimiK3Mode::Idle, + }; + } + + fn apply_event(&mut self, event: KimiK3Event, output: &mut UnifiedParserOutput) -> Result<()> { + match event { + KimiK3Event::Text { len } => output.push_text(self.buffer[..len].to_string()), + KimiK3Event::Reasoning { len } => { + output.push_reasoning(self.buffer[..len].to_string()); + } + KimiK3Event::Skip => {} + KimiK3Event::ThinkOpen => self.mode = KimiK3Mode::Reasoning, + KimiK3Event::ThinkClose => self.mode = KimiK3Mode::Idle, + KimiK3Event::ResponseOpen => self.mode = KimiK3Mode::Response, + KimiK3Event::ResponseClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::ToolsOpen => self.mode = KimiK3Mode::Tools, + KimiK3Event::ToolsClose => self.mode = KimiK3Mode::Epilogue, + KimiK3Event::MessageEnd => self.mode = KimiK3Mode::Done, + KimiK3Event::CallOpen { name, index } => { + self.mode = KimiK3Mode::Call { + name, + index, + scan: MarkerScanState::default(), + }; + } + KimiK3Event::CallComplete { arguments } => { + let mode = std::mem::replace(&mut self.mode, KimiK3Mode::Tools); + let KimiK3Mode::Call { name, index, .. } = mode else { + return Err(parsing_failed!( + "Kimi K3 call completion without an active tool call" + )); + }; + // An empty/garbage call block without a tool name is dropped, + // matching the Python parser. + if name.is_empty() { + return Ok(()); + } + + let tool_index = self.call_ids.len(); + self.call_ids.push(tool_call_id_for(&name, index.as_deref())); + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments, + }); + } + } + Ok(()) + } + + fn reset_state(&mut self) -> String { + self.mode = KimiK3Mode::Idle; + self.call_ids.clear(); + std::mem::take(&mut self.buffer) + } +} + +impl UnifiedParser for KimiK3UnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.call_ids.clear(); + self.initialize_mode(prompt_token_ids); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> { + Some(&KIMI_K3_STRUCTURAL_TAG_BUILDER) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.call_ids.get(tool_index).map(String::as_str) + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + parse_next_kimi_k3_event(input, &mut self.mode) + })? { + self.apply_event(event, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + KimiK3Mode::Idle | KimiK3Mode::Response => { + output.push_text(std::mem::take(&mut self.buffer)); + } + KimiK3Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + KimiK3Mode::Epilogue | KimiK3Mode::Done => self.buffer.clear(), + // A tools channel truncated between complete calls loses only its + // closing markers; keep the calls already emitted. + KimiK3Mode::Tools if self.buffer.is_empty() => {} + KimiK3Mode::Tools | KimiK3Mode::Call { .. } => { + return Err(parsing_failed!("incomplete Kimi K3 tool call")); + } + } + + // Keep call_ids so tool_call_id() stays available after the stream ends. + self.mode = KimiK3Mode::Idle; + Ok(output) + } + + fn reset(&mut self) -> String { + self.reset_state() + } +} + +/// Build the API-side tool-call ID from the XTML one-based `index` attribute. +/// +/// The ID uses the zero-based call ordinal; XTML's message index stays +/// one-based when rendering tool result messages. +fn tool_call_id_for(name: &str, index: Option<&str>) -> String { + match index { + None => name.to_string(), + Some(raw) => match raw.parse::() { + Ok(one_based) => format!("{name}:{}", one_based - 1), + Err(_) => format!("{name}:{raw}"), + }, + } +} + +/// Parse one Kimi K3 event from buffered streaming input. +fn parse_next_kimi_k3_event( + input: &mut KimiK3Input<'_>, + mode: &mut KimiK3Mode, +) -> ModalResult { + match mode { + KimiK3Mode::Idle => parse_idle_event(input), + KimiK3Mode::Reasoning => parse_reasoning_event(input), + KimiK3Mode::Response => parse_response_event(input), + KimiK3Mode::Epilogue => parse_epilogue_event(input), + KimiK3Mode::Tools => parse_tools_event(input), + KimiK3Mode::Call { scan, .. } => call_body_event(input, scan), + KimiK3Mode::Done => parse_done_event(input), + } +} + +/// Parse an event while waiting for the next channel open. +fn parse_idle_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_OPEN).value(KimiK3Event::ThinkOpen), + literal(RESPONSE_OPEN).value(KimiK3Event::ResponseOpen), + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_idle_text_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `think` channel. +fn parse_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(THINK_CLOSE).value(KimiK3Event::ThinkClose), + // `<|end_of_msg|>` can reach the parser under `ignore_eos` or + // `include_stop_str_in_output`; never leak it into reasoning. + literal(END_OF_MSG).value(KimiK3Event::MessageEnd), + safe_reasoning_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `response` channel. +fn parse_response_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(RESPONSE_CLOSE).value(KimiK3Event::ResponseClose), + // The response body also implicitly ends at a `tools` channel. + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + safe_response_text_event, + )) + .parse_next(input) +} + +/// Parse an event after the response channel closed. +fn parse_epilogue_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + literal(TOOLS_OPEN).value(KimiK3Event::ToolsOpen), + message_end_event, + skip_epilogue_noise_event, + )) + .parse_next(input) +} + +/// Parse an event inside the `tools` channel, between `call` blocks. +fn parse_tools_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt(( + call_open_event, + literal(TOOLS_CLOSE).value(KimiK3Event::ToolsClose), + // Defensive: an unterminated tools channel still ends with the message. + message_end_event, + skip_tools_noise_event, + )) + .parse_next(input) +} + +/// Parse a message close or end-of-message marker. +fn message_end_event(input: &mut KimiK3Input<'_>) -> ModalResult { + alt((literal(MESSAGE_CLOSE), literal(END_OF_MSG))) + .value(KimiK3Event::MessageEnd) + .parse_next(input) +} + +/// Ignore everything after the assistant message closed. +fn parse_done_event(input: &mut KimiK3Input<'_>) -> ModalResult { + rest.value(KimiK3Event::Skip).parse_next(input) +} + +/// Parse safe text while waiting for the next channel marker. +fn safe_idle_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Parse safe reasoning before the think close marker. +fn safe_reasoning_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, REASONING_MARKERS).map(|len| KimiK3Event::Reasoning { len }) +} + +/// Parse safe response text before the next channel marker. +fn safe_response_text_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, RESPONSE_MARKERS).map(|len| KimiK3Event::Text { len }) +} + +/// Skip non-content noise after the response channel closed. +fn skip_epilogue_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, EPILOGUE_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Skip non-content noise between `call` blocks. +fn skip_tools_noise_event(input: &mut KimiK3Input<'_>) -> ModalResult { + safe_text_len_mul(input, TOOLS_MARKERS).map(|_| KimiK3Event::Skip) +} + +/// Parse a `call` open tag into its tool name and one-based index. +fn call_open_event(input: &mut KimiK3Input<'_>) -> ModalResult { + let (attrs,) = seq!( + _: literal(CALL_OPEN), + take_until(0.., SEP), + _: literal(SEP), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + Ok(KimiK3Event::CallOpen { + name: attr_value(&attrs, "tool").unwrap_or_default().to_string(), + index: attr_value(&attrs, "index") + .filter(|index| !index.is_empty()) + .map(str::to_string), + }) +} + +/// Parse the buffered call body through `<|close|>call<|sep|>` into a +/// completed call. +fn call_body_event( + input: &mut KimiK3Input<'_>, + scan: &mut MarkerScanState, +) -> ModalResult { + let (body,) = seq!( + take_until_marker(CALL_CLOSE, scan), + _: literal(CALL_CLOSE), + ) + .parse_next(input)?; + let arguments = parse_call_arguments(body)?; + + Ok(KimiK3Event::CallComplete { arguments }) +} + +/// Parse a complete `call` body into the OpenAI-style arguments JSON string. +/// +/// The body is either one raw `json` block (passed through unmodified) or a +/// sequence of typed `argument` blocks (converted per their `type` tags). +fn parse_call_arguments(body: &str) -> ModalResult { + let mut input = body; + terminated( + delimited(ws0, alt((json_block_arguments, typed_arguments)), ws0), + eof, + ) + .parse_next(&mut input) + .map_err(|_| xtml_error("Kimi K3 call body")) +} + +/// Parse one raw `json` argument block, passing its body through unmodified. +fn json_block_arguments(input: &mut &str) -> ModalResult { + let (raw,) = seq!( + _: literal(JSON_OPEN), + _: take_until(0.., SEP), // attrs (`type="object"`), unused on decode + _: literal(SEP), + take_until(0.., JSON_CLOSE), + _: literal(JSON_CLOSE), + ) + .parse_next(input)?; + Ok(raw.to_string()) +} + +/// Parse typed `argument` blocks into a serialized JSON object, preserving +/// argument order. +fn typed_arguments(input: &mut &str) -> ModalResult { + let pairs: Vec<(String, Value)> = + repeat(0.., terminated(argument_block, ws0)).parse_next(input)?; + let arguments = pairs.into_iter().collect::>(); + serde_json::to_string(&arguments).map_err(|_| xtml_error("Kimi K3 arguments")) +} + +/// Parse one typed `argument` block into its key/value pair. +fn argument_block(input: &mut &str) -> ModalResult<(String, Value)> { + let (attrs, raw) = seq!( + _: literal(ARG_OPEN), + take_until(0.., SEP), + _: literal(SEP), + take_until(0.., ARG_CLOSE), + _: literal(ARG_CLOSE), + ) + .parse_next(input)?; + let attrs = parse_tag_attrs(attrs)?; + + let key = attr_value(&attrs, "key").unwrap_or_default().to_string(); + let arg_type = attr_value(&attrs, "type").unwrap_or("string"); + Ok((key, decode_argument_value(arg_type, raw))) +} + +/// Decode one typed argument value per its XTML `type` tag. +/// +/// `string` values pass the raw text through (the renderer emits them +/// unescaped); other types are JSON-decoded, falling back to the raw text on +/// malformed payloads so one quirky value does not fail the whole call. +fn decode_argument_value(arg_type: &str, raw: &str) -> Value { + if arg_type == "string" { + return Value::String(raw.to_string()); + } + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string())) +} + +/// Parse a complete XTML attribute string like ` tool="get_weather" index="1"`. +fn parse_tag_attrs(attrs: &str) -> ModalResult> { + let mut input = attrs; + terminated(repeat(0.., preceded(ws1, tag_attr)), (ws0, eof)) + .parse_next(&mut input) + .map_err(|_| xtml_error("XTML tag attributes")) +} + +/// Parse one XTML `key="value"` attribute pair. +fn tag_attr(input: &mut &str) -> ModalResult<(String, String)> { + seq!( + take_while(1.., |char: char| char.is_alphanumeric() || char == '_').map(str::to_string), + _: literal("=\""), + take_till(0.., '"').map(unescape_attr_value), + _: literal("\""), + ) + .parse_next(input) +} + +/// Reverse XTML attribute escaping: `"` first, then `&` (the inverse +/// of the encode order). +fn unescape_attr_value(value: &str) -> String { + value.replace(""", "\"").replace("&", "&") +} + +/// Look up one parsed attribute value by key. +fn attr_value<'a>(attrs: &'a [(String, String)], key: &str) -> Option<&'a str> { + attrs.iter().find(|(name, _)| name == key).map(|(_, value)| value.as_str()) +} + +/// Build a cut error for determinably malformed XTML structure. +fn xtml_error(label: &'static str) -> ErrMode { + let mut error = ContextError::new(); + error.push(StrContext::Label(label)); + ErrMode::Cut(error) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer as _; + use vllm_tokenizer::test_utils::TestTokenizer; + + use super::{ + END_OF_MSG, KimiK3UnifiedParser, OPEN, RESPONSE_CLOSE, RESPONSE_OPEN, SEP, THINK_CLOSE, + THINK_OPEN, TOOLS_CLOSE, TOOLS_OPEN, + }; + use crate::tool::ToolCallDelta; + use crate::unified::{ + UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, + }; + + const OPEN_ID: u32 = 256; + const CLOSE_ID: u32 = 257; + const SEP_ID: u32 = 258; + const END_OF_MSG_ID: u32 = 259; + + fn tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_special_token(OPEN, OPEN_ID) + .with_special_token("<|close|>", CLOSE_ID) + .with_special_token(SEP, SEP_ID) + .with_special_token(END_OF_MSG, END_OF_MSG_ID) + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for KimiK3UnifiedParser { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_parser() -> KimiK3UnifiedParser { + KimiK3UnifiedParser::new(&[], Arc::new(tokenizer())).unwrap() + } + + fn collect_stream(parser: &mut KimiK3UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + /// Split `text` into small chunks to stress marker-split handling. + fn char_chunks(text: &str, size: usize) -> Vec { + let chars: Vec = text.chars().collect(); + chars.chunks(size).map(|chunk| chunk.iter().collect()).collect() + } + + fn arg(key: &str, arg_type: &str, value: &str) -> String { + format!( + "{OPEN}argument key=\"{key}\" type=\"{arg_type}\"{SEP}{value}<|close|>argument{SEP}" + ) + } + + fn call(attrs: &str, body: &str) -> String { + format!("{OPEN}call {attrs}{SEP}{body}<|close|>call{SEP}") + } + + fn thinking_output(reasoning: &str, response: &str, tools_body: &str) -> String { + let mut output = format!("{THINK_OPEN}{reasoning}{THINK_CLOSE}"); + output.push_str(&format!("{RESPONSE_OPEN}{response}{RESPONSE_CLOSE}")); + if !tools_body.is_empty() { + output.push_str(&format!("{TOOLS_OPEN}{tools_body}{TOOLS_CLOSE}")); + } + output.push_str("<|close|>message<|sep|>"); + output + } + + fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta { + output.calls().first().expect("expected one tool call").clone() + } + + #[test] + fn kimi_k3_create_requires_structural_tokens() { + let error = match KimiK3UnifiedParser::new(&[], Arc::new(TestTokenizer::new())) { + Ok(_) => panic!("expected missing token error"), + Err(error) => error, + }; + + assert!(matches!( + error, + UnifiedParserError::MissingToken { token } if token == OPEN + )); + } + + #[test] + fn kimi_k3_parses_reasoning_response_and_typed_tool_call() { + let body = [ + arg("city", "string", "Hangzhou"), + arg("days", "number", "1.5"), + arg("detailed", "boolean", "true"), + arg("filters", "object", r#"{"kind":"rain"}"#), + arg("hours", "array", "[8,20]"), + ] + .concat(); + let text = thinking_output( + "Need the weather tool.", + "I'll check.", + &call("tool=\"get_weather\" index=\"1\"", &body), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.reasoning_text(), "Need the weather tool."); + assert_eq!(output.normal_text(), "I'll check."); + let call = first_call(&output); + assert_eq!(call.tool_index, 0); + assert_eq!(call.name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&call.arguments).unwrap(), + json!({ + "city": "Hangzhou", + "days": 1.5, + "detailed": true, + "filters": { "kind": "rain" }, + "hours": [8, 20], + }) + ); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + } + + #[test] + fn kimi_k3_arguments_preserve_order_and_number_formatting() { + let body = [ + arg("y", "number", "1.0"), + arg("x", "number", "2"), + arg("items", "array", r#"["left","right"]"#), + ] + .concat(); + let text = thinking_output("t", "", &call("tool=\"add\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + first_call(&output).arguments, + r#"{"y":1.0,"x":2,"items":["left","right"]}"# + ); + } + + #[test] + fn kimi_k3_streaming_splits_markers_across_chunks() { + let text = thinking_output( + "step by step", + "the answer", + &call("tool=\"calc\" index=\"1\"", &arg("x", "number", "42")), + ); + + for size in [1, 3, 7] { + let chunks = char_chunks(&text, size); + let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); + let output = collect_stream(&mut test_parser(), &chunk_refs); + + assert_eq!(output.reasoning_text(), "step by step", "chunk size {size}"); + assert_eq!(output.normal_text(), "the answer", "chunk size {size}"); + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":42}"#); + } + } + + #[test] + fn kimi_k3_streaming_emits_text_incrementally() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let first = parser.parse_chunk("Hel").unwrap(); + assert_eq!(first.normal_text(), "Hel"); + + let second = parser.parse_chunk("lo<|close|>resp").unwrap(); + assert_eq!(second.normal_text(), "lo"); + + let mut output = parser.parse_chunk("onse<|sep|>").unwrap(); + output.append(parser.finish().unwrap()); + assert_eq!(output.normal_text(), ""); + } + + #[test] + fn kimi_k3_initialize_think_prefill_starts_in_reasoning() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>think<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!( + "reasoning{THINK_CLOSE}{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}" + )) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reasoning"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_initialize_response_prefill_starts_in_response() { + let mut parser = test_parser(); + let prompt = tokenizer() + .encode( + "<|open|>message role=\"assistant\"<|sep|><|open|>response<|sep|>", + false, + ) + .unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("answer{RESPONSE_CLOSE}<|close|>message{SEP}")) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn kimi_k3_initialize_message_open_prefill_starts_idle() { + let mut parser = test_parser(); + let prompt = + tokenizer().encode("<|open|>message role=\"assistant\"<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + + let output = parser + .parse_complete(&format!("{THINK_OPEN}reason{THINK_CLOSE}{RESPONSE_OPEN}hi")) + .unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "hi"); + } + + #[test] + fn kimi_k3_plain_text_falls_through_as_text() { + let output = collect_stream(&mut test_parser(), &["plain ", "answer"]); + + assert_eq!(output.normal_text(), "plain answer"); + assert!(output.reasoning_text().is_empty()); + assert!(output.calls().is_empty()); + } + + #[test] + fn kimi_k3_tool_call_waits_for_close_marker() { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + + let argument = arg("x", "number", "1"); + for chunk in [ + TOOLS_OPEN, + "<|open|>call tool=\"calc\" index=\"1\"<|sep|>", + argument.as_str(), + ] { + output.append(parser.parse_chunk(chunk).unwrap()); + assert!(output.calls().is_empty()); + } + + output.append(parser.parse_chunk("<|close|>call<|sep|>").unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + assert_eq!(first_call(&output).arguments, r#"{"x":1}"#); + } + + #[test] + fn kimi_k3_parses_multiple_tool_calls() { + let tools_body = format!( + "{}{}", + call( + "tool=\"get_weather\" index=\"1\"", + &arg("city", "string", "SF") + ), + call("tool=\"get_time\" index=\"2\"", ""), + ); + let text = thinking_output("t", "r", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(calls[0].arguments, r#"{"city":"SF"}"#); + assert_eq!(calls[1].tool_index, 1); + assert_eq!(calls[1].name.as_deref(), Some("get_time")); + assert_eq!(calls[1].arguments, "{}"); + assert_eq!(parser.tool_call_id(0), Some("get_weather:0")); + assert_eq!(parser.tool_call_id(1), Some("get_time:1")); + } + + #[test] + fn kimi_k3_json_block_arguments_pass_through_raw() { + // Spacing and key order must survive unmodified: raw `json` blocks are + // not validated or normalized. + let raw = r#"{"b": 1, "a": [2 , 3]}"#; + let body = format!("{OPEN}json type=\"object\"{SEP}{raw}<|close|>json{SEP}"); + let text = thinking_output("t", "", &call("tool=\"run\" index=\"1\"", &body)); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, raw); + } + + #[test] + fn kimi_k3_string_argument_passes_raw_text_through() { + let value = "line one\nline two {\"not\": \"json\"} & "; + let text = thinking_output( + "t", + "", + &call( + "tool=\"write\" index=\"1\"", + &arg("content", "string", value), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!( + serde_json::from_str::(&first_call(&output).arguments).unwrap(), + json!({ "content": value }) + ); + } + + #[test] + fn kimi_k3_malformed_typed_argument_falls_back_to_raw_text() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"calc\" index=\"1\"", + &arg("x", "number", "not a number"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).arguments, r#"{"x":"not a number"}"#); + } + + #[test] + fn kimi_k3_attribute_values_are_unescaped() { + let text = thinking_output( + "t", + "", + &call( + "tool=\"a"b&c\" index=\"1\"", + &arg("key", "string", "value"), + ), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(first_call(&output).name.as_deref(), Some("a\"b&c")); + } + + #[test] + fn kimi_k3_call_without_tool_name_is_dropped() { + let tools_body = format!( + "{}{}", + call("index=\"1\"", &arg("x", "number", "1")), + call("tool=\"real\" index=\"2\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + let calls = output.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[0].name.as_deref(), Some("real")); + assert_eq!(parser.tool_call_id(0), Some("real:1")); + } + + #[test] + fn kimi_k3_tool_call_id_follows_index_attribute() { + let tools_body = format!( + "{}{}{}", + call("tool=\"first\" index=\"3\"", ""), + call("tool=\"second\"", ""), + call("tool=\"third\" index=\"x\"", ""), + ); + let text = thinking_output("t", "", &tools_body); + + let mut parser = test_parser(); + parser.parse_complete(&text).unwrap(); + + assert_eq!(parser.tool_call_id(0), Some("first:2")); + assert_eq!(parser.tool_call_id(1), Some("second")); + assert_eq!(parser.tool_call_id(2), Some("third:x")); + } + + #[test] + fn kimi_k3_ignores_output_after_message_close() { + let mut parser = test_parser(); + let output = parser + .parse_complete(&format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}<|close|>message{SEP}junk{END_OF_MSG}" + )) + .unwrap(); + + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn kimi_k3_epilogue_noise_is_not_content() { + let text = format!( + "{RESPONSE_OPEN}answer{RESPONSE_CLOSE}\n{TOOLS_OPEN}{}{TOOLS_CLOSE}\n<|close|>message{SEP}", + call("tool=\"calc\" index=\"1\"", ""), + ); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_finish_flushes_unclosed_reasoning() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk(&format!("{THINK_OPEN}still thinking")).unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.reasoning_text(), "still thinking"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn kimi_k3_finish_flushes_partial_marker_as_text() { + let mut parser = test_parser(); + let mut output = parser.parse_chunk("answer<|clo").unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(output.normal_text(), "answer<|clo"); + } + + #[test] + fn kimi_k3_finish_fails_mid_tool_call() { + let mut parser = test_parser(); + parser + .parse_chunk(&format!( + "{TOOLS_OPEN}<|open|>call tool=\"calc\" index=\"1\"<|sep|>{}", + arg("x", "number", "1") + )) + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert!(error.to_report_string().contains("incomplete Kimi K3 tool call")); + } + + #[test] + fn kimi_k3_finish_after_truncated_tools_keeps_complete_calls() { + let mut parser = test_parser(); + let mut output = parser + .parse_chunk(&format!( + "{TOOLS_OPEN}{}", + call("tool=\"calc\" index=\"1\"", &arg("x", "number", "1")) + )) + .unwrap(); + output.append(parser.finish().unwrap()); + + assert_eq!(first_call(&output).name.as_deref(), Some("calc")); + } + + #[test] + fn kimi_k3_malformed_call_attributes_fail_fast() { + let mut parser = test_parser(); + let error = parser + .parse_chunk(&format!("{TOOLS_OPEN}<|open|>call garbage attrs<|sep|>")) + .unwrap_err(); + + assert!(error.to_report_string().contains("XTML tag attributes")); + } + + #[test] + fn kimi_k3_empty_response_channel_emits_nothing() { + let text = thinking_output("t", "", &call("tool=\"calc\" index=\"1\"", "")); + + let mut parser = test_parser(); + let output = parser.parse_complete(&text).unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.reasoning_text(), "t"); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn kimi_k3_reset_returns_buffered_text() { + let mut parser = test_parser(); + let prompt = tokenizer().encode("<|open|>response<|sep|>", false).unwrap(); + parser.initialize(&prompt).unwrap(); + parser.parse_chunk("answer<|close|>resp").unwrap(); + + let raw = parser.reset(); + + assert_eq!(raw, "<|close|>resp"); + } +} diff --git a/rust/src/parser/src/unified/kimi_k3/structural_tag.rs b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs new file mode 100644 index 00000000000..d3515ad2dc3 --- /dev/null +++ b/rust/src/parser/src/unified/kimi_k3/structural_tag.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Structural-tag grammar for Kimi K3 XTML tool calls. + +use serde_json::{Map, Value}; +use xgrammar_structural_tag::Result; +use xgrammar_structural_tag::builders::{ + StructuralTagBuilder, StructuralTagContext, StructuralTagOptions, +}; +use xgrammar_structural_tag::format::{Format, JsonSchemaFormat, StructuralTag, TagFormat}; +use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam, function_parameters}; + +use super::{ + ARG_CLOSE, CALL_CLOSE, END_OF_MSG, JSON_CLOSE, JSON_OPEN, MESSAGE_CLOSE, OPEN, RESPONSE_CLOSE, + RESPONSE_OPEN, SEP, THINK_CLOSE, TOOLS_CLOSE, TOOLS_OPEN, +}; + +pub(super) static KIMI_K3_STRUCTURAL_TAG_BUILDER: KimiK3StructuralTagBuilder = + KimiK3StructuralTagBuilder; + +const XTML_TYPES: &[&str] = &["string", "number", "boolean", "null", "object", "array"]; + +/// Kimi K3 XTML structural-tag builder. +#[derive(Debug, Clone, Copy, Default)] +pub struct KimiK3StructuralTagBuilder; + +impl StructuralTagBuilder for KimiK3StructuralTagBuilder { + fn build(&self, ctx: StructuralTagContext<'_>) -> Result { + let mut elements = response_prefix(ctx.options.reasoning); + + let tools = match ctx.tool_choice { + // Serving lowering filters empty tools before calling the builder, + // while direct `build_structural_tag(..., auto, ...)` calls do not. + BuilderToolChoice::Auto if ctx.function_tools.is_empty() => None, + BuilderToolChoice::Auto => Some(Format::optional(tools_channel( + ctx.function_tools, + ctx.options, + ))), + BuilderToolChoice::Forced | BuilderToolChoice::Required => { + Some(tools_channel(ctx.function_tools, ctx.options)) + } + }; + if let Some(tools) = tools { + elements.push(tools); + } + elements.push(Format::optional(Format::const_string(MESSAGE_CLOSE))); + + Ok(StructuralTag::new(Format::sequence(elements))) + } +} + +fn response_prefix(reasoning: bool) -> Vec { + let mut elements = Vec::new(); + if reasoning { + elements.push(Format::tag( + "", + Format::any_text_excluding(&[THINK_CLOSE, END_OF_MSG]), + THINK_CLOSE, + )); + elements.push(Format::const_string(RESPONSE_OPEN)); + } else { + elements.push(Format::optional(Format::const_string(RESPONSE_OPEN))); + } + elements.push(Format::tag( + "", + Format::any_text_excluding(&[RESPONSE_CLOSE, TOOLS_OPEN, MESSAGE_CLOSE, END_OF_MSG]), + RESPONSE_CLOSE, + )); + elements +} + +fn tools_channel(tools: &[FunctionToolParam], options: StructuralTagOptions) -> Format { + let calls = tools.iter().map(|tool| call_tag(tool, options)).collect(); + Format::tag( + TOOLS_OPEN, + Format::tags_with_separator(calls, "", true, false), + TOOLS_CLOSE, + ) +} + +fn call_tag(tool: &FunctionToolParam, options: StructuralTagOptions) -> TagFormat { + let parameters = function_parameters(&tool.function); + let call_body = Format::or(vec![ + typed_arguments(¶meters, options), + raw_json_arguments(¶meters, options), + ]); + + TagFormat::new( + format!( + "{OPEN}call tool=\"{}\" index=\"", + escape_attr_value(&tool.function.name) + ), + Format::sequence(vec![ + Format::regex("[1-9][0-9]*"), + Format::const_string(format!("\"{SEP}")), + call_body, + ]), + CALL_CLOSE, + ) +} + +fn typed_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + let Some(schema) = parameters.as_object() else { + return if parameters == &Value::Bool(false) { + Format::const_string("") + } else { + Format::star(permissive_argument()) + }; + }; + let Some(properties) = schema.get("properties").and_then(Value::as_object) else { + return Format::star(permissive_argument()); + }; + if properties.is_empty() { + return Format::star(permissive_argument()); + } + + let root_defs = root_definitions(schema); + let arguments = properties + .iter() + .flat_map(|(key, schema)| argument_tags(key, schema, &root_defs, options)) + .map(Format::Tag) + .collect::>(); + let arguments = match arguments.as_slice() { + [argument] => argument.clone(), + _ => Format::or(arguments), + }; + // Keep typed arguments order-agnostic and non-unique, but do not allow an + // empty call when the root schema declares required properties. + if schema + .get("required") + .and_then(Value::as_array) + .is_some_and(|required| !required.is_empty()) + { + Format::plus(arguments) + } else { + Format::star(arguments) + } +} + +fn argument_tags( + key: &str, + schema: &Value, + root_defs: &Map, + options: StructuralTagOptions, +) -> Vec { + let types = schema_types(schema); + types + .into_iter() + .map(|xtml_type| { + let content = if xtml_type == "string" { + string_argument_content(schema) + } else { + json_schema( + attach_root_definitions(&narrow_schema_type(schema, xtml_type), root_defs), + options, + ) + }; + TagFormat::new( + format!( + "{OPEN}argument key=\"{}\" type=\"{xtml_type}\"{SEP}", + escape_attr_value(key) + ), + content, + ARG_CLOSE, + ) + }) + .collect() +} + +fn schema_types(schema: &Value) -> Vec<&'static str> { + let Some(schema) = schema.as_object() else { + return XTML_TYPES.to_vec(); + }; + let mut types = Vec::new(); + match schema.get("type") { + Some(Value::String(value)) => push_schema_type(&mut types, value), + Some(Value::Array(values)) => { + for value in values.iter().filter_map(Value::as_str) { + push_schema_type(&mut types, value); + } + } + _ => {} + } + if types.is_empty() + && let Some(value) = schema.get("const") + { + push_value_type(&mut types, value); + } + if types.is_empty() + && let Some(values) = schema.get("enum").and_then(Value::as_array) + { + for value in values { + push_value_type(&mut types, value); + } + } + if types.is_empty() { + XTML_TYPES.to_vec() + } else { + types + } +} + +fn push_value_type(types: &mut Vec<&'static str>, value: &Value) { + let xtml_type = match value { + Value::String(_) => "string", + Value::Number(_) => "number", + Value::Bool(_) => "boolean", + Value::Null => "null", + Value::Object(_) => "object", + Value::Array(_) => "array", + }; + if !types.contains(&xtml_type) { + types.push(xtml_type); + } +} + +fn narrow_schema_type(schema: &Value, xtml_type: &str) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + let json_type = if xtml_type == "number" && explicitly_integer_only(&schema) { + "integer" + } else { + xtml_type + }; + schema.insert("type".to_string(), Value::String(json_type.to_string())); + Value::Object(schema) +} + +fn explicitly_integer_only(schema: &Map) -> bool { + match schema.get("type") { + Some(Value::String(value)) => value == "integer", + Some(Value::Array(values)) => { + let values = values.iter().filter_map(Value::as_str).collect::>(); + values.contains(&"integer") && !values.contains(&"number") + } + _ => false, + } +} + +fn push_schema_type(types: &mut Vec<&'static str>, json_type: &str) { + let xtml_type = match json_type { + "string" => Some("string"), + "integer" | "number" => Some("number"), + "boolean" => Some("boolean"), + "null" => Some("null"), + "object" => Some("object"), + "array" => Some("array"), + _ => None, + }; + if let Some(xtml_type) = xtml_type + && !types.contains(&xtml_type) + { + types.push(xtml_type); + } +} + +fn string_argument_content(schema: &Value) -> Format { + let Some(schema) = schema.as_object() else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + let values = schema + .get("enum") + .and_then(Value::as_array) + .cloned() + .or_else(|| schema.get("const").cloned().map(|value| vec![value])); + let Some(values) = values else { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + }; + if values.is_empty() + || values.len() > 256 + || values + .iter() + .any(|value| value.as_str().is_none_or(|value| value.contains("<|"))) + { + return Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]); + } + let values = values.iter().filter_map(Value::as_str).collect::>(); + match values.as_slice() { + [value] => Format::const_string(*value), + _ => Format::or(values.into_iter().map(Format::const_string).collect()), + } +} + +fn raw_json_arguments(parameters: &Value, options: StructuralTagOptions) -> Format { + Format::tag( + format!("{JSON_OPEN} type=\"object\"{SEP}"), + json_schema(parameters.clone(), options), + JSON_CLOSE, + ) +} + +fn permissive_argument() -> Format { + let key = Format::regex(r#"(?:[^<\"&]|&(?:amp|quot);|<[^|])*"#); + let alternatives = XTML_TYPES + .iter() + .map(|xtml_type| { + Format::sequence(vec![ + key.clone(), + Format::const_string(format!("\" type=\"{xtml_type}\"{SEP}")), + if *xtml_type == "string" { + Format::any_text_excluding(&[ARG_CLOSE, CALL_CLOSE]) + } else { + Format::json_schema(Value::Bool(true)) + }, + ]) + }) + .collect(); + Format::tag( + format!("{OPEN}argument key=\""), + Format::or(alternatives), + ARG_CLOSE, + ) +} + +fn json_schema(schema: Value, options: StructuralTagOptions) -> Format { + Format::JsonSchema( + JsonSchemaFormat::new(schema) + .with_any_order(options.any_order) + .with_max_whitespace_cnt(options.max_whitespace_cnt), + ) +} + +fn root_definitions(schema: &Map) -> Map { + ["$defs", "definitions"] + .into_iter() + .filter_map(|key| schema.get(key).map(|value| (key.to_string(), value.clone()))) + .collect() +} + +fn attach_root_definitions(schema: &Value, root_defs: &Map) -> Value { + let Some(mut schema) = schema.as_object().cloned() else { + return schema.clone(); + }; + for (key, value) in root_defs { + schema.entry(key.clone()).or_insert_with(|| value.clone()); + } + Value::Object(schema) +} + +fn escape_attr_value(value: &str) -> String { + value.replace('&', "&").replace('"', """) +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use serde_json::json; + use xgrammar_structural_tag::builders::StructuralTagOptions; + use xgrammar_structural_tag::{ + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag, + }; + + use super::KimiK3StructuralTagBuilder; + + fn tool(name: &str, parameters: serde_json::Value) -> ToolParam { + ToolParam::Function(FunctionToolParam::new( + FunctionDefinition::new(name).with_parameters(parameters), + )) + } + + #[test] + fn required_structural_tag_matches_xtml_channels() { + let tools = vec![tool( + "get_weather", + json!({ + "$defs": { + "place": { "type": "object", "properties": { "city": { "type": "string" } } } + }, + "type": "object", + "properties": { + "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] }, + "place": { "$ref": "#/$defs/place", "type": "object" } + }, + "required": ["place"] + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap(); + + expect![[r##"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"optional","content":{"type":"const_string","value":"<|open|>response<|sep|>"}},{"type":"tag","begin":"","content":{"type":"any_text","excludes":["<|close|>response<|sep|>","<|open|>tools<|sep|>","<|close|>message<|sep|>","<|end_of_msg|>"]},"end":"<|close|>response<|sep|>"},{"type":"tag","begin":"<|open|>tools<|sep|>","content":{"type":"tags_with_separator","tags":[{"begin":"<|open|>call tool=\"get_weather\" index=\"","content":{"type":"sequence","elements":[{"type":"regex","pattern":"[1-9][0-9]*"},{"type":"const_string","value":"\"<|sep|>"},{"type":"or","elements":[{"type":"plus","content":{"type":"or","elements":[{"type":"tag","begin":"<|open|>argument key=\"unit\" type=\"string\"<|sep|>","content":{"type":"or","elements":[{"type":"const_string","value":"celsius"},{"type":"const_string","value":"fahrenheit"}]},"end":"<|close|>argument<|sep|>"},{"type":"tag","begin":"<|open|>argument key=\"place\" type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$ref":"#/$defs/place","type":"object","$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}}},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>argument<|sep|>"}]}},{"type":"tag","begin":"<|open|>json type=\"object\"<|sep|>","content":{"type":"json_schema","json_schema":{"$defs":{"place":{"type":"object","properties":{"city":{"type":"string"}}}},"type":"object","properties":{"unit":{"type":"string","enum":["celsius","fahrenheit"]},"place":{"$ref":"#/$defs/place","type":"object"}},"required":["place"]},"style":"json","any_order":false,"max_whitespace_cnt":null},"end":"<|close|>json<|sep|>"}]}]},"end":"<|close|>call<|sep|>"}],"separator":"","at_least_one":true,"stop_after_first":false},"end":"<|close|>tools<|sep|>"},{"type":"optional","content":{"type":"const_string","value":"<|close|>message<|sep|>"}}]}}"##]].assert_eq(&tag.to_json_string().unwrap()); + } + + #[test] + fn typed_arguments_require_one_tag_only_for_nonempty_required() { + let required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query"] + }), + StructuralTagOptions::default(), + ); + let optional = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } } + }), + StructuralTagOptions::default(), + ); + let empty_required = super::typed_arguments( + &json!({ + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": [] + }), + StructuralTagOptions::default(), + ); + + assert_eq!(serde_json::to_value(required).unwrap()["type"], "plus"); + assert_eq!(serde_json::to_value(optional).unwrap()["type"], "star"); + assert_eq!( + serde_json::to_value(empty_required).unwrap()["type"], + "star" + ); + } + + #[test] + fn reasoning_grammar_starts_inside_prefilled_think_channel() { + let tools = vec![tool("ping", json!({ "type": "object", "properties": {} }))]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::auto(), + StructuralTagOptions::default().with_reasoning(true), + ) + .unwrap(); + let value = serde_json::to_value(tag).unwrap(); + + assert_eq!( + value["format"]["elements"][0]["end"], + "<|close|>think<|sep|>" + ); + assert_eq!( + value["format"]["elements"][1]["value"], + "<|open|>response<|sep|>" + ); + assert_eq!(value["format"]["elements"][3]["type"], "optional"); + } + + #[test] + fn forced_choice_keeps_only_the_named_tool() { + let tools = vec![ + tool("search", json!({ "type": "object" })), + tool("lookup", json!({ "type": "object" })), + ]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::function("lookup"), + StructuralTagOptions::default().with_reasoning(false), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains("lookup")); + assert!(!tag.contains("search")); + } + + #[test] + fn union_argument_content_matches_its_xtml_type() { + let tools = vec![tool( + "set_count", + json!({ + "type": "object", + "properties": { + "count": { "type": ["integer", "null"] } + } + }), + )]; + let tag = build_structural_tag( + KimiK3StructuralTagBuilder, + &tools, + ToolChoice::required(), + StructuralTagOptions::default(), + ) + .unwrap() + .to_json_string() + .unwrap(); + + assert!(tag.contains(r#"type=\"number\""#)); + assert!(tag.contains(r#""json_schema":{"type":"integer"}"#), "{tag}"); + assert!(tag.contains(r#"type=\"null\""#)); + assert!(tag.contains(r#""json_schema":{"type":"null"}"#), "{tag}"); + } + + #[test] + fn unsafe_string_enum_falls_back_as_a_whole() { + let format = super::string_argument_content(&json!({ + "type": "string", + "enum": ["safe", "<|unsafe"] + })); + + assert_eq!(serde_json::to_value(format).unwrap()["type"], "any_text"); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 4700211ae93..ae5524759df 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -6,10 +6,12 @@ mod combined; mod gemma4; mod inkling; +mod kimi_k3; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; pub use inkling::InklingUnifiedParser; +pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser}; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index d75f3fab162..2c4050f8f1e 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -89,6 +89,15 @@ pub(super) fn prepare_chat_request( )?; let template_kwargs = request.chat_template_kwargs.unwrap_or_default(); + let response_format = + request.response_format.as_ref().map(serde_json::to_value).transpose().map_err( + |error| { + ApiError::invalid_request( + format!("failed to serialize response_format: {error}"), + Some("response_format"), + ) + }, + )?; let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) @@ -150,6 +159,7 @@ pub(super) fn prepare_chat_request( generation_prompt_mode, chat_template: request.chat_template, reasoning_effort: request.reasoning_effort, + response_format, template_kwargs, }, tools: convert_tools(request.tools)?, @@ -400,6 +410,7 @@ fn convert_tool_choice(tool_choice: Option<&ToolChoice>) -> Result &'static str { let model_type = config.effective_model_type(); match model_type { - Some("kimi" | "kimi_k2" | "kimi_k25" | "deepseek_v3") => KIMI_PATTERN, + Some("kimi" | "kimi_k2" | "kimi_k25" | "kimi_k3" | "deepseek_v3") => KIMI_PATTERN, _ => CL100K_BASE_PATTERN, } } @@ -817,6 +817,7 @@ mod tests { #[test] fn tiktoken_detects_kimi_pattern_from_model_type() { let kimi = config_json!({ "model_type": "kimi_k25" }); + let kimi_k3 = config_json!({ "model_type": "kimi_k3" }); let baseten_kimi = config_json!({ "model_type": "deepseek_v3" }); let nested_kimi = config_json!({ "model_type": "composite_wrapper", @@ -830,6 +831,7 @@ mod tests { let missing = config_json!({ "text_config": {} }); assert_eq!(detect_bpe_pattern(&kimi), KIMI_PATTERN); + assert_eq!(detect_bpe_pattern(&kimi_k3), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&baseten_kimi), KIMI_PATTERN); assert_eq!(detect_bpe_pattern(&nested_kimi), CL100K_BASE_PATTERN); assert_eq!(detect_bpe_pattern(&generic), CL100K_BASE_PATTERN); From bb3b61f2fd2333ab165ebaba13f133db4210b9f2 Mon Sep 17 00:00:00 2001 From: Julien Debache Date: Tue, 28 Jul 2026 23:57:22 +0200 Subject: [PATCH 45/67] perf: dispatch non-grouped bias-less topk routing methods to fused path (#49618) Signed-off-by: jdebache --- tests/kernels/moe/test_routing.py | 116 ++++++++++++++++++ .../layers/fused_moe/router/router_factory.py | 68 ++++++---- 2 files changed, 160 insertions(+), 24 deletions(-) diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 62a4968a0d1..9b12fef0454 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -8,9 +8,19 @@ import torch from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + FusedTopKBiasRouter, +) +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( + FusedTopKRouter, +) +from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + GroupedTopKRouter, +) from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) @@ -36,6 +46,112 @@ TOP_KS = [2, 4, 6] NUM_EXPERTS = [8, 16, 64] +def test_degenerate_grouped_config_uses_standard_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, FusedTopKRouter) + hidden_states, router_logits = make_test_data(32, 256, 128) + + topk_weights, topk_ids = router.select_experts(hidden_states, router_logits) + baseline_weights, baseline_ids = baseline_fused_topk( + router_logits, + top_k=4, + renormalize=True, + ) + + assert_routing_results_close( + topk_weights, + topk_ids, + baseline_weights, + baseline_ids, + ) + + +def test_multiple_expert_groups_use_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=8, + topk_group=4, + scoring_func="softmax", + renormalize=True, + ) + + assert isinstance(router, GroupedTopKRouter) + + +def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + + +def test_degenerate_grouped_config_with_bias_keeps_routed_scale() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, FusedTopKBiasRouter) + assert router.routed_scaling_factor == 1.1 + + +def test_degenerate_deepseek_v3_routing_stays_grouped() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + renormalize=True, + e_score_correction_bias=torch.empty(128), + ) + + assert isinstance(router, GroupedTopKRouter) + assert router.routing_method_type == RoutingMethodType.DeepSeekV3 + + +def test_single_expert_group_with_non_unit_scale_uses_grouped_topk() -> None: + router = create_fused_moe_router( + top_k=4, + global_num_experts=128, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + renormalize=True, + routed_scaling_factor=1.1, + ) + + assert isinstance(router, GroupedTopKRouter) + + def setup_eplb_state( enable_eplb: bool, global_num_experts: int ) -> EplbLayerState | None: diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index c7cfccbe64b..ff4874b20e6 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -9,6 +9,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, + get_routing_method_type, ) from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501 AiterSharedRoutedFusedMoERouter, @@ -67,7 +68,8 @@ def create_fused_moe_router( The selection logic follows this priority order: 1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set 2. ZeroExpertRouter - if zero_expert_type is not None - 3. GroupedTopKRouter - if use_grouped_topk is True + 3. GroupedTopKRouter - if use_grouped_topk is True and the grouping is not + degenerate (at most one group, with topk_group <= 1) 4. CustomRoutingRouter - if custom_routing_function is not None 5. FusedTopKBiasRouter - if e_score_correction_bias is not None 6. AiterSharedRoutedFusedMoERouter - if num_fused_shared_experts > 0 @@ -143,30 +145,48 @@ def create_fused_moe_router( "num_expert_group and topk_group must be provided when " "use_grouped_topk is True" ) - grouped_topk_router = GroupedTopKRouter( - top_k=top_k, - global_num_experts=global_num_experts, - eplb_state=eplb_state, - num_expert_group=num_expert_group, - topk_group=topk_group, - renormalize=renormalize, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - num_fused_shared_experts=num_fused_shared_experts, - ) - if ( - grouped_topk_router.routing_method_type != RoutingMethodType.Unspecified - or num_expert_group > 1 - or topk_group > 1 - ): - return grouped_topk_router - # If routing_method for GroupedTopKRouter is Unspecified and there is only - # one group, fallback to standard top-k routing - use_grouped_topk = False - num_expert_group = None - topk_group = None + # For topk_group <= 1, grouped implementation is pure overhead. + degenerate_grouping = num_expert_group <= 1 and topk_group <= 1 + # FusedTopKRouter cannot apply routed_scaling_factor, FusedTopKBiasRouter can. + scaling_handled_downstream = ( + routed_scaling_factor == 1.0 or e_score_correction_bias is not None + ) + + # Degenerating must not change the advertised routing method, which drives + # kernel selection. num_expert_group only affects it for biased routing. + def advertised_routing_method(groups: int | None) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=scoring_func, + top_k=top_k, + renormalize=renormalize, + num_expert_group=groups, + has_e_score_bias=e_score_correction_bias is not None, + routed_scaling_factor=routed_scaling_factor, + ) + + routing_method_preserved = advertised_routing_method( + num_expert_group + ) == advertised_routing_method(None) + + if not ( + degenerate_grouping + and scaling_handled_downstream + and routing_method_preserved + ): + return GroupedTopKRouter( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=num_fused_shared_experts, + ) + # Otherwise fall through to the non-grouped chain below. if custom_routing_function is not None: return CustomRoutingRouter( From a07fac758f2cae58f6ebad29da726e9432012c13 Mon Sep 17 00:00:00 2001 From: Ruinan Ma <97484148+mrn3088@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:28:28 -0700 Subject: [PATCH 46/67] [Perf] Zero-copy torch.Tensor pickling in shm_broadcast MessageQueue (#48442) Signed-off-by: Ruinan Ma Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- tests/distributed/test_shm_broadcast.py | 149 ++++++++++++++++++ tools/pre_commit/check_forbidden_imports.py | 1 + .../device_communicators/shm_broadcast.py | 86 +++++++++- 3 files changed, 233 insertions(+), 3 deletions(-) diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index 17957924051..0b413965032 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io +import pickle import random import threading import time @@ -10,12 +12,15 @@ from unittest import mock import multiprocess as mp import numpy as np import pytest +import torch import torch.distributed as dist from vllm.distributed.device_communicators import shm_broadcast from vllm.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, + _rebuild_tensor, + _reduce_tensor, check_shm_free_space, ) from vllm.distributed.utils import StatelessProcessGroup @@ -354,6 +359,150 @@ def test_message_queue_busy_to_idle(): distributed_run(worker_fn_test_busy_to_idle, 4) +@worker_fn_wrapper +def worker_fn_tensor_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + message_queue = MessageQueue.create_from_process_group( + dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank + ) + + # Both ranks construct the identical reference payload. + torch.manual_seed(42) + payload = { + # 2MiB: rides the shm ring as an out-of-band buffer (the receiving + # side must copy out of the reusable ring chunk). + "mid": torch.randn(1024, 512), + # 16MiB > max_chunk_bytes: overflows to the zmq socket (the + # receiving side aliases the zmq.Frame zero-copy). + "big": torch.randn(4096, 2048, dtype=torch.bfloat16), + "nested": ["plain", 123, {"inner": torch.arange(5)}], + } + + if rank == writer_rank: + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + assert wrapped_reduce.call_count == 3 + # Cycle the ring (max_chunks=4) several times over so that aliased + # ring chunks would be overwritten. + for i in range(16): + message_queue.enqueue({"junk": torch.full((1024, 512), float(i))}) + else: + received = message_queue.dequeue(timeout=30) + for key in ("mid", "big"): + assert torch.equal(received[key], payload[key]), key + assert received[key].dtype == payload[key].dtype, key + assert torch.equal(received["nested"][2]["inner"], torch.arange(5)) + + snapshot = received["mid"].clone() + for i in range(16): + junk = message_queue.dequeue(timeout=30) + assert torch.equal(junk["junk"], torch.full((1024, 512), float(i))) + # Tensors received via the shm ring must not alias chunk memory + # that the writer has reused for subsequent messages. + assert torch.equal(received["mid"], snapshot) + # Rebuilt tensors must be writable, like regular tensors. + received["mid"] += 1.0 + received["big"][0, 0] = 1.0 + + dist.barrier() + print(f"tensor broadcast passed the test! Rank {rank}") + + +def test_tensor_broadcast(): + distributed_run(worker_fn_tensor_broadcast, 2) + + +def _dumps_oob(obj) -> tuple[bytes, list]: + """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor + dispatch table + out-of-band buffers >= 1MiB.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + bio = io.BytesIO() + pickler = pickle.Pickler( + bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback + ) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + pickler.dump(obj) + return bio.getvalue(), buffers + + +@pytest.mark.parametrize( + "case", + [ + "small", + "mid", + "bf16", + "fp8", + "empty", + "scalar", + "noncontig", + "requires_grad", + "conj", + "param", + ], +) +def test_tensor_pickle_roundtrip(case: str): + tensor = { + # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers. + "small": lambda: torch.randn(100, 10), + "mid": lambda: torch.randn(1024, 512), + # Dtypes numpy doesn't recognize. + "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16), + "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn), + # Shape edge cases. + "empty": lambda: torch.empty(0, 8), + "scalar": lambda: torch.tensor(3.14), + "noncontig": lambda: torch.randn(64, 64).t(), + # These fall back to torch's default reducer. + "requires_grad": lambda: torch.randn(8, 8, requires_grad=True), + "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(), + "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False), + }[case]() + + data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))}) + received = pickle.loads(data, buffers=buffers)["tensor"] + + assert received.shape == tensor.shape + assert received.dtype == tensor.dtype + if tensor.dtype == torch.float8_e4m3fn: + assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8)) + else: + assert torch.equal(received, tensor) + assert received.requires_grad == tensor.requires_grad + assert isinstance(received, type(tensor)) + if tensor.numel() and not tensor.requires_grad: + # Rebuilt tensors must be writable, like regular tensors. + received.view(-1)[0] = 1.0 + + +@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"]) +def test_reduce_tensor_fallback(case: str): + """Tensors the zero-copy reducer can't safely alias must fall back to + torch's default reduction.""" + if case == "cuda": + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + tensor = torch.randn(4, device="cuda") + elif case == "requires_grad": + tensor = torch.randn(8, requires_grad=True) + else: + tensor = torch.randn(4, dtype=torch.complex64).conj() + + reduced = _reduce_tensor(tensor) + assert reduced[0] is not _rebuild_tensor + + @pytest.mark.parametrize("should_warn", [False, True]) def test_reader_timeout_caps_indefinite_waits(should_warn): with ( diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index a788cecc6ce..52a95ce1d8d 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -48,6 +48,7 @@ CHECK_IMPORTS = { "vllm/distributed/device_communicators/shm_object_storage.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", + "tests/distributed/test_shm_broadcast.py", "tests/distributed/test_weight_transfer.py", "vllm/utils/hashing.py", "tests/multimodal/media/test_base.py", diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index afabdf18c80..e59b14f7a6b 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copyreg import functools +import io import os import pickle import shutil @@ -385,6 +387,70 @@ class ShmRingBuffer: yield buf +def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor: + """Rebuild a tensor from an out-of-band pickle buffer. + + Counterpart of `_reduce_tensor`. Note that pickle passes the original + buffer-providing object from `loads(buffers=...)` straight to this + function (no `PickleBuffer` wrapper on the receiving side), so `buf` is + a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` + if the buffer was serialized in-band. + """ + dtype = getattr(torch, dtype_str) + assert isinstance(dtype, torch.dtype) + if isinstance(buf, zmq.Frame): + # ZMQ frames own their message memory independently of any context, + # so the tensor can safely alias it with zero copies. The tensor's + # storage keeps the frame (and thus its bytes) alive via a strong + # reference for as long as the tensor is. + try: + return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape) + except ValueError: + # Empty or read-only frame buffer; fall through to the copy path. + pass + # Shared-memory ring buffer chunks are reused by the writer once all + # readers have marked them read, so we must copy out of them. bytearray + # (vs bytes) keeps the resulting tensor writable, matching normal tensor + # semantics. + raw = bytearray(buf) + if not raw: + assert 0 in shape + return torch.empty(shape, dtype=dtype) + return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape) + + +def _reduce_tensor(tensor: torch.Tensor): + """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling. + + `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle + byte stream via `torch.serialization` and never emits a `PickleBuffer`, + which defeats the out-of-band buffer handling in `MessageQueue.enqueue`. + This reducer instead exposes the tensor's memory directly, so large + tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue + without being copied into and back out of the pickled message. + """ + if ( + tensor.device.type == "cpu" + and tensor.layout == torch.strided + and not tensor.requires_grad + ): + try: + # The uint8 view exposes the raw bytes via the buffer protocol, + # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...). + # reshape(-1) first so that 0-dim tensors can be viewed as well. + raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy() + except RuntimeError: + # Exotic tensors (e.g. with the conjugate bit set) that don't + # support aliasing views; let torch handle them. + pass + else: + dtype_str = str(tensor.dtype).removeprefix("torch.") + return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str) + + # Fall back to torch's default (copying) reduction. + return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL) + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -771,9 +837,23 @@ class MessageQueue: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps( - obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback - ) + # CPU tensors are routed through `_reduce_tensor` so that their + # bytes are emitted as out-of-band buffers instead of being + # copied into the pickle stream by torch's default reducer. + # Start from `copyreg.dispatch_table` to preserve globally + # registered reducers (e.g. `re.Pattern`); the per-pickler + # dispatch table would otherwise shadow them. + dispatch_table = dict(copyreg.dispatch_table) + dispatch_table[torch.Tensor] = _reduce_tensor + with io.BytesIO() as bio: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) + pickler.dispatch_table = dispatch_table + pickler.dump(obj) + all_buffers[0] = bio.getvalue() if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: From e7f6a39db8d1f05ea923b5f976141db838b11c44 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Wed, 29 Jul 2026 07:47:48 +0800 Subject: [PATCH 47/67] [Test] Make EPD correctness tests configurable for XPU (#50110) Signed-off-by: zhenwei-intel --- .../integration/run_epd_correctness_test.sh | 85 ++++++++++++------- .../integration/test_epd_correctness.py | 10 ++- 2 files changed, 64 insertions(+), 31 deletions(-) diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index 65716444a57..c58df0c076a 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh @@ -21,12 +21,19 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # Model to test MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}" # Set 1 to use multimodal prompts; else to use text-only USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}" -MM_FLAG="" -if [ "$USE_MM_PROMPTS" = "1" ]; then - MM_FLAG="--use_mm_prompts" +USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}" +TEST_FLAGS=() +if [[ "$USE_MM_PROMPTS" == "1" ]]; then + TEST_FLAGS+=(--use_mm_prompts) +fi +if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then + TEST_FLAGS+=(--skip_two_image_prompt) fi # GPU configuration @@ -36,6 +43,16 @@ GPU_D="${GPU_D:-2}" GPU_SINGLE="${GPU_SINGLE:-$GPU_P}" GPU_PD="${GPU_PD:-$GPU_P}" +# Device platform and affinity environment variable +DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}" +if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then + if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then + DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK" + else + DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES" + fi +fi + # Port ENCODE_PORT="${ENCODE_PORT:-19534}" PREFILL_PORT="${PREFILL_PORT:-19535}" @@ -87,11 +104,12 @@ run_baseline() { # Start baseline instance echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT" - CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" vllm serve "$MODEL" \ --port "$PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ - --max-num-seqs 128 \ + --gpu-memory-utilization 0.9 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ > "$LOG_PATH"/baseline.log 2>&1 & @@ -112,7 +130,7 @@ run_baseline() { --model_name "$MODEL" \ --mode baseline \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup baseline echo "Stopping baseline instance..." @@ -139,14 +157,15 @@ run_epd_1e_1pd() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -160,12 +179,13 @@ run_epd_1e_1pd() { # Start prefill+decode instance echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_PD" vllm serve "$MODEL" \ --port "$PREFILL_DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -212,7 +232,7 @@ run_epd_1e_1pd() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1PD Correctness Test finished" @@ -242,14 +262,15 @@ run_baseline_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ vllm serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -260,14 +281,15 @@ run_baseline_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ vllm serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -309,7 +331,7 @@ run_baseline_1p_1d() { --model_name "$MODEL" \ --mode baseline_pd \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "Stopping PD (1P+1D) instances..." @@ -339,14 +361,15 @@ run_epd_1e_1p_1d() { # Start encoder instance echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \ + env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \ --port "$ENCODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ --gpu-memory-utilization 0.01 \ --enable-request-id-headers \ --no-enable-prefix-caching \ --max-num-batched-tokens 114688 \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -360,14 +383,15 @@ run_epd_1e_1p_1d() { # Start prefill instance echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT" - CUDA_VISIBLE_DEVICES="$GPU_P" \ + env "$DEVICE_AFFINITY_ENV=$GPU_P" \ VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \ vllm serve "$MODEL" \ --port "$PREFILL_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --ec-transfer-config '{ "ec_connector": "ECExampleConnector", @@ -385,14 +409,15 @@ run_epd_1e_1p_1d() { # Start decode instance echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT" - CUDA_VISIBLE_DEVICES="$GPU_D" \ + env "$DEVICE_AFFINITY_ENV=$GPU_D" \ VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \ vllm serve "$MODEL" \ --port "$DECODE_PORT" \ + --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ - --gpu-memory-utilization 0.7 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --enable-request-id-headers \ - --max-num-seqs 128 \ + --max-num-seqs "$MAX_NUM_SEQS" \ --allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \ --kv-transfer-config '{ "kv_connector": "NixlConnector", @@ -438,7 +463,7 @@ run_epd_1e_1p_1d() { --model_name "$MODEL" \ --mode disagg \ --baseline_file "$BASELINE_PD_FILE" \ - $MM_FLAG + "${TEST_FLAGS[@]}" # Cleanup echo "✓✓ 1E+1P+1D Correctness Test finished" @@ -465,7 +490,7 @@ run_epd_1e_1pd # Step 3: Test baseline 1P + 1D run_baseline_1p_1d -# Step 4: Test 1E + 1P + 1D +# # Step 4: Test 1E + 1P + 1D run_epd_1e_1p_1d # Cleanup output file diff --git a/tests/v1/ec_connector/integration/test_epd_correctness.py b/tests/v1/ec_connector/integration/test_epd_correctness.py index eae4b742724..ece73efa42d 100644 --- a/tests/v1/ec_connector/integration/test_epd_correctness.py +++ b/tests/v1/ec_connector/integration/test_epd_correctness.py @@ -192,6 +192,12 @@ def main(): help="Use multimodal prompts (default: use text-only for quick testing)", ) + parser.add_argument( + "--skip_two_image_prompt", + action="store_true", + help="Skip the two-image multimodal prompt", + ) + args = parser.parse_args() print(f"Service URL: {args.service_url}") @@ -221,7 +227,9 @@ def main(): # Select prompts to use if args.use_mm_prompts: - test_prompts = SAMPLE_PROMPTS_MM + test_prompts = ( + SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM + ) print("Using multimodal prompts") else: test_prompts = SAMPLE_PROMPTS_TEXT From 5369f7b7b89e48d396fa98a9d33f3fba654af8ee Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Wed, 29 Jul 2026 01:57:20 +0200 Subject: [PATCH 48/67] [MXFP8][ROCm] Fix MXFP8 MoE backend selection (#49747) Signed-off-by: Felix Marty --- .../moe/test_mxfp8_aiter_backend_selection.py | 44 ++++- tests/models/quantization/test_mxfp8.py | 166 ++++++++++++++++++ .../fused_moe/experts/aiter_mxfp8_moe.py | 27 +++ .../fused_moe/experts/mxfp8_emulation_moe.py | 8 +- .../fused_moe/experts/mxfp8_native_moe.py | 9 +- .../layers/fused_moe/oracle/mxfp8.py | 51 +----- 6 files changed, 244 insertions(+), 61 deletions(-) diff --git a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py index 7c2fdbabe29..aa1747f4153 100644 --- a/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py +++ b/tests/kernels/moe/test_mxfp8_aiter_backend_selection.py @@ -19,9 +19,17 @@ if not current_platform.is_rocm(): pytest.skip("This test can only run on ROCm.", allow_module_level=True) from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402 +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + MoEActivation, +) from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402 + _AITER_SWIGLU_ALPHA, + _AITER_SWIGLU_BETA, AiterMxfp8Experts, ) +from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( # noqa: E402 + Mxfp8NativeTritonExperts, +) from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402 FusedMoEActivationFormat, ) @@ -33,6 +41,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402 _SUPPORTED_BACKENDS, _mxfp8_backend_to_kernel_cls, _select_kernel_cls, + select_mxfp8_moe_backend, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 kMxfp8Dynamic, @@ -43,7 +52,17 @@ _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe" def _config(ep_size: int = 1): - cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144) + # AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and + # alpha/beta so is_supported_config doesn't reject the config on those grounds. + cfg = make_dummy_moe_config( + num_experts=128, + experts_per_token=4, + hidden_dim=6144, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + cfg = dataclasses.replace( + cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA + ) if ep_size != 1: cfg = dataclasses.replace( cfg, @@ -76,12 +95,6 @@ def test_aiter_mxfp8_registered(): ] -def test_triton_selectable(): - assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8 - # Not auto-selected (only reachable explicitly), so FlyDSL still wins auto. - assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS - - @pytest.mark.parametrize("ep_size", [1, 2]) def test_ep_supported(ep_size): """FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask.""" @@ -133,3 +146,20 @@ def test_explicit_moe_backend_aiter(): pytest.raises(ValueError, match="flydsl package"), ): _select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1)) + + +def test_gfx950_picks_aiter(): + """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins.""" + with _flydsl_installed(True): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is AiterMxfp8Experts + + +def test_gfx942_picks_triton(): + """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton + dot_scaled backend wins instead.""" + with _flydsl_installed(False): + backend, experts_cls = select_mxfp8_moe_backend(_config()) + assert backend is Fp8MoeBackend.TRITON_MXFP8 + assert experts_cls is Mxfp8NativeTritonExperts diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py index 7c250d11576..c12a72a09c0 100644 --- a/tests/models/quantization/test_mxfp8.py +++ b/tests/models/quantization/test_mxfp8.py @@ -17,8 +17,10 @@ diverse prompts from ``tests/prompts/example.txt``. """ import pytest +import torch from tests.quantization.utils import is_quant_method_supported +from vllm.platforms import current_platform from ..utils import check_logprobs_close @@ -81,6 +83,170 @@ def test_mxfp8_logprobs( ) +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_activation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_requires_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + ) + + with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"): + select_mxfp8_moe_backend(config) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="AITER MXFP8 MoE backend is ROCm-only.", +) +@pytest.mark.quant_model +def test_mxfp8_aiter_accepts_swigluoai_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe + from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, + ) + + monkeypatch.setattr( + aiter_mxfp8_moe.AiterMxfp8Experts, + "_supports_current_device", + staticmethod(lambda: True), + ) + monkeypatch.setattr( + aiter_mxfp8_moe, + "is_aiter_mxfp8_moe_available", + lambda: True, + ) + + config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + in_dtype=torch.bfloat16, + device="cuda", + routing_method=RoutingMethodType.Renormalize, + moe_backend="aiter", + swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA, + swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA, + ) + + backend, experts_cls = select_mxfp8_moe_backend(config) + + assert backend == Fp8MoeBackend.AITER_MXFP8 + assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts + + @pytest.mark.skipif( not is_quant_method_supported("mxfp8"), reason="mxfp8 is not supported on this GPU type (requires sm_100+).", diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c5330f3b438..59115df13b5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -6,10 +6,13 @@ ``convert_to_fp8_moe_kernel_format``. """ +import math + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8TritonExpertsBase, ) @@ -17,6 +20,9 @@ from vllm.platforms import current_platform logger = init_logger(__name__) +_AITER_SWIGLU_ALPHA = 1.702 +_AITER_SWIGLU_BETA = 1.0 + def is_aiter_mxfp8_moe_available() -> bool: """True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl`` @@ -93,6 +99,27 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): return False, ( "kernel requires the aiter flydsl package, which is not installed" ) + if ( + is_supported + and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE + ): + return False, ( + "kernel hardcodes SwiGLU-OAI activation and requires " + f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; " + f"got activation={moe_config.activation.value}" + ) + if is_supported and ( + moe_config.swiglu_alpha is None + or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA) + or moe_config.swiglu_beta is None + or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA) + ): + return False, ( + "kernel hardcodes SwiGLU-OAI with " + f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; " + f"got swiglu_alpha={moe_config.swiglu_alpha} and " + f"swiglu_beta={moe_config.swiglu_beta}" + ) return is_supported, reason def apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py index 71dd7634a69..ad6083251fb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py @@ -107,17 +107,13 @@ class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase): limit = self.quant_config.gemm1_clamp_limit if limit is None: raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit") - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) apply_moe_activation( activation, output, input, clamp_limit=float(limit), - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, ) return super().activation(activation, output, input) diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index 9839756880a..e8c7dc96921 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -362,10 +362,7 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase): expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, ): - alpha = self.quant_config.gemm1_alpha - alpha = 1.702 if alpha is None else float(alpha) - beta = self.quant_config.gemm1_beta - beta = 1.0 if beta is None else float(beta) + # `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`. limit = self.quant_config.gemm1_clamp_limit limit = None if limit is None else float(limit) out = fused_moe_mxfp8_native( @@ -376,8 +373,8 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase): self.w2_scale_val, topk_weights, topk_ids, - alpha=alpha, - beta=beta, + alpha=self.gemm1_alpha, + beta=self.gemm1_beta, limit=limit, global_num_experts=global_num_experts, expert_map=expert_map, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index b9086cfa48a..2a03eacea8e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -12,10 +12,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp8Dynamic, kMxfp8Static, ) -from vllm.platforms import current_platform logger = init_logger(__name__) +# Ordered by priority. _SUPPORTED_BACKENDS = ( Fp8MoeBackend.FLASHINFER_TRTLLM, Fp8MoeBackend.DEEPGEMM, @@ -26,6 +26,8 @@ _SUPPORTED_BACKENDS = ( # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, Fp8MoeBackend.HUMMING, + Fp8MoeBackend.TRITON_MXFP8, + Fp8MoeBackend.EMULATION, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls( return [AiterMxfp8Experts] if backend == Fp8MoeBackend.TRITON_MXFP8: - # Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e. - # dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise. - # Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``. - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + Mxfp8NativeTritonExperts, + ) - return [Mxfp8NativeTritonExperts] + return [Mxfp8NativeTritonExperts] + if backend == Fp8MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8EmulationTritonExperts, ) @@ -105,35 +104,6 @@ def _select_kernel_cls( ) -def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: - """ROCm fallback when no auto-selected MXFP8 backend is available. - - The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by - ``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or - explicitly via ``--moe-backend aiter``; this fallback handles the rest - (native dot_scaled on gfx950, else BF16 emulation). - """ - - if current_platform.supports_mx(): - from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( - Mxfp8NativeTritonExperts, - ) - - logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.") - return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts - - from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( - Mxfp8EmulationTritonExperts, - ) - - logger.info_once( - "No native MXFP8 MoE backend available on this device; " - "MXFP8 weights will be dequantized to BF16 once at load time and the " - "MoE will run in BF16 (no per-step dequant)." - ) - return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts - - def select_mxfp8_moe_backend( config: FusedMoEConfig, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: @@ -167,8 +137,5 @@ def select_mxfp8_moe_backend( logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) return backend, experts_cls - # simplify the logic for rocm, refactor later when more backends are supported - if current_platform.is_rocm(): - return _select_rocm_mxfp8_backend() - + # TODO: add debug log with reason. raise ValueError("No MXFP8 MoE backends available.") From 176256b9628d2d8d22db41d7a72c15540f438026 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 28 Jul 2026 18:58:10 -0500 Subject: [PATCH 49/67] [ROCm][CI] Stabilize ROCm audio streaming test (#50163) Signed-off-by: Andreas Karatzas --- .../multimodal/openai/chat_completion/test_audio.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py index fa0f141afee..a1c13d9e339 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_audio.py @@ -10,6 +10,7 @@ import pytest_asyncio from tests.utils import RemoteOpenAIServer from vllm.assets.audio import AudioAsset from vllm.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio +from vllm.platforms import current_platform MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b" TEST_AUDIO_URLS = [ @@ -18,6 +19,10 @@ TEST_AUDIO_URLS = [ ] MAXIMUM_AUDIOS = 2 +# Disable prefix caching on ROCm to reduce non-determinism in +# streaming-vs-non-streaming comparisons. +_ROCM_ARGS = ["--no-enable-prefix-caching"] if current_platform.is_rocm() else [] + @pytest.fixture(scope="module") def server(): @@ -32,6 +37,7 @@ def server(): "--trust-remote-code", "--limit-mm-per-prompt", json.dumps({"audio": MAXIMUM_AUDIOS}), + *_ROCM_ARGS, ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: From fe65aa6a97981c815800daa4bf49d8356b8c4988 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:12:50 -0400 Subject: [PATCH 50/67] [CI][NIXL] Fix flaky DP+EP test port conflict (#50171) Signed-off-by: Divakar Verma --- tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index 22682e02bea..fc5c04a1ad0 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -227,9 +227,12 @@ run_tests_for_model() { # Calculate side channel port SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE)) INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE)) - DECODER_INTERNAL_PORT_ENV= + # For non-DP mode, set VLLM_PORT to pin the internal port; + # For DP mode, set VLLM_DP_MASTER_PORT instead to avoid race condition. if [[ -z "${DP_EP:-}" ]]; then DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT" + else + DECODER_INTERNAL_PORT_ENV="VLLM_DP_MASTER_PORT=$INTERNAL_PORT" fi echo "Starting decode instance $i on GPU $GPU_ID, port $PORT" From 56f31af62afe6553369b14af225bfafb788e99a0 Mon Sep 17 00:00:00 2001 From: Kevin Glynn Date: Tue, 28 Jul 2026 20:17:35 -0400 Subject: [PATCH 51/67] [Bugfix] Fix /wake_up crash on hybrid models (Mamba/DeltaNet) (#41602) Signed-off-by: Kevin Glynn --- tests/v1/worker/test_gpu_model_runner.py | 56 ++++++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 13 ++++-- 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 79e3a60e981..d2d89a74fce 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1669,3 +1669,59 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): with pytest.raises(ValueError, match="max_num_seqs"): runner.initialize_kv_cache(kv_cache_config) + + +class TestInitFp8KvScalesHybridModels: + """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries. + + Hybrid models (Mamba, DeltaNet) store per-layer state as a list of tensors + rather than a single tensor. init_fp8_kv_scales must iterate both forms. + """ + + @staticmethod + def _make_runner_stub(kv_caches): + runner = Mock(spec=GPUModelRunner) + runner.cache_config = SimpleNamespace(cache_dtype="fp8_e4m3") + runner.kv_caches = kv_caches + runner.compilation_config = SimpleNamespace(static_forward_context={}) + runner.init_fp8_kv_scales = GPUModelRunner.init_fp8_kv_scales.__get__( + runner, GPUModelRunner + ) + return runner + + def test_zeroes_both_tensor_and_list_entries(self): + single_tensor = torch.ones(4, 8) + list_tensors = [torch.ones(2, 4), torch.ones(3, 6)] + + runner = self._make_runner_stub([single_tensor, list_tensors]) + runner.init_fp8_kv_scales() + + assert (single_tensor == 0).all() + assert all((t == 0).all() for t in list_tensors) + + def test_skips_none_entries(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([None, tensor, None]) + runner.init_fp8_kv_scales() + + assert (tensor == 0).all() + + def test_noop_when_kv_cache_not_quantized(self): + tensor = torch.ones(4, 8) + runner = self._make_runner_stub([tensor]) + runner.cache_config.cache_dtype = "auto" + runner.init_fp8_kv_scales() + + assert (tensor == 1).all() + + def test_mixed_none_tensor_and_list(self): + t1 = torch.ones(2, 2) + t2 = torch.ones(3, 3) + list_entry = [torch.ones(1, 1), torch.ones(1, 1)] + + runner = self._make_runner_stub([None, t1, list_entry, None, t2]) + runner.init_fp8_kv_scales() + + assert (t1 == 0).all() + assert (t2 == 0).all() + assert all((t == 0).all() for t in list_entry) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index dd31ddab672..0859b6dacd8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -991,9 +991,16 @@ class GPUModelRunner( return kv_caches = getattr(self, "kv_caches", []) - for cache_tensor in kv_caches: - if cache_tensor is not None: - cache_tensor.zero_() + for cache_entry in kv_caches: + if cache_entry is None: + continue + # Hybrid models (Mamba, DeltaNet) store per-layer state as a + # list of tensors rather than a single tensor. + if isinstance(cache_entry, list): + for t in cache_entry: + t.zero_() + else: + cache_entry.zero_() k_attr_names = ("_k_scale", "k_scale") v_attr_names = ("_v_scale", "v_scale") From 6fbbcf2151c34cd9313c223b8cbf19afa77f2ea6 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 18:09:41 -0700 Subject: [PATCH 52/67] [BugFix] Stop dummy runs from writing mamba state through stale block-table rows (#49757) Signed-off-by: Nick Hill Signed-off-by: Jeff Ma Co-authored-by: Jeff Ma --- tests/v1/worker/test_gpu_block_table.py | 63 +++++++++++++++++++++++++ vllm/v1/worker/block_table.py | 5 ++ vllm/v1/worker/gpu/block_table.py | 8 +++- vllm/v1/worker/gpu/model_runner.py | 4 ++ vllm/v1/worker/gpu_model_runner.py | 8 +++- 5 files changed, 86 insertions(+), 2 deletions(-) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 31acd475ade..365e365a574 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -130,3 +130,66 @@ def test_block_tables_apply_staged_writes_single_group(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) + + +def test_v1_block_table_move_row_clears_vacated_row(): + """condense() moves the last row into a freed slot; the vacated row must + not keep stale block ids. Padded dummy-run batches dereference stale rows + as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode + padding) and write state in place there — corrupting the blocks' new + owner once they are reallocated, e.g. to an in-flight NIXL load.""" + from vllm.v1.worker.block_table import BlockTable + + block_table = BlockTable( + block_size=16, + max_num_reqs=4, + max_num_blocks_per_req=8, + max_num_batched_tokens=64, + pin_memory=False, + device=torch.device("cuda"), + kernel_block_size=16, + cp_kv_cache_interleave_size=1, + ) + block_table.add_row([7, 8, 9], row_idx=0) + block_table.add_row([4, 5], row_idx=1) + + block_table.move_row(1, 0) + + assert block_table.block_table.np[0, :2].tolist() == [4, 5] + assert block_table.num_blocks_per_row[0] == 2 + # The vacated source row routes to the reserved null block. + assert block_table.num_blocks_per_row[1] == 0 + assert (block_table.block_table.np[1] == 0).all() + + +def test_get_dummy_block_tables_returns_zeroed_rows(): + """Dummy runs bypass the gather, so the persistent input_block_tables + hold the previous real step's rows. Mamba/GDN metadata routes in-place + state writes through block_table[:, 0] (dummy slot mappings are + PAD-filled, state indices are not), so stale rows would direct dummy + state writes at freed — possibly reallocated — blocks. + get_dummy_block_tables must hand out zeroed (null block) rows while + preserving the persistent storage address for CUDA graphs.""" + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=4, + max_num_batched_tokens=64, + max_num_blocks_per_group=[8], + device=device, + kernel_block_sizes=[16], + ) + # Simulate a real step: stage a request's blocks and gather them into + # the persistent input block tables. + block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) + torch.accelerator.synchronize() + assert block_tables.input_block_tables[0][0, 0].item() == 1 + + dummy = block_tables.get_dummy_block_tables(num_reqs=1) + torch.accelerator.synchronize() + assert (dummy[0] == 0).all() + # CUDA graph invariant: same persistent tensor, not a fresh allocation. + assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 332eda4cbdf..4d9f256ee00 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -144,6 +144,11 @@ class BlockTable: block_table_np = self.block_table.np block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks] self.num_blocks_per_row[tgt] = num_blocks + # Clear the vacated source row: dummy-run batches dereference stale + # rows as mamba state slots and write state in place there, possibly + # after the blocks have been freed and reallocated. + block_table_np[src, :num_blocks] = 0 + self.num_blocks_per_row[src] = 0 def swap_row(self, src: int, tgt: int) -> None: src_tgt, tgt_src = [src, tgt], [tgt, src] diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index df5d2be629d..5d17868b174 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -163,7 +163,13 @@ class BlockTables: # Therefore, this method must return the persistent tensor # with the same memory address as that used during the model's forward pass, # rather than allocating a new tensor. - return tuple(block_table[:num_reqs] for block_table in self.input_block_tables) + # + # Zero the rows so dummy runs write mamba state to the reserved null + # block rather than through the previous real step's (stale) block + # ids, which may point at blocks since freed and reallocated. + return tuple( + block_table[:num_reqs].zero_() for block_table in self.input_block_tables + ) def compute_slot_mappings( self, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9182ada4394..b16e75bf45e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1290,6 +1290,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings, self.attn_groups, self.kv_cache_config, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) input_ids = input_batch.input_ids diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0859b6dacd8..3019b7750aa 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6064,7 +6064,13 @@ class GPUModelRunner( num_reqs=num_reqs_padded, max_query_len=max_query_len, ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), - for_cudagraph_capture=is_graph_capturing, + # FULL replay reads capture-time metadata buffers. Re-stage them + # from the zeroed dummy block tables instead of retaining state + # indices from the previous real batch. + for_cudagraph_capture=( + is_graph_capturing + or cudagraph_runtime_mode == CUDAGraphMode.FULL + ), slot_mappings=slot_mappings_by_group, use_spec_decode=self.speculative_config is not None, ) From 7f4c52f2ba38a77d5341fb90a31c2cedfc0066b8 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 28 Jul 2026 18:41:48 -0700 Subject: [PATCH 53/67] [CI] Add comment-based Buildkite triggers (#50132) Signed-off-by: khluu Co-authored-by: OpenAI Codex --- .github/workflows/new_pr_bot.yml | 4 +- .github/workflows/pre-commit.yml | 2 +- .github/workflows/run-ci-command.yml | 40 ++ .github/workflows/scripts/run_ci_command.py | 580 ++++++++++++++++++ .../workflows/scripts/test_run_ci_command.py | 362 +++++++++++ docs/contributing/README.md | 6 +- 6 files changed, 989 insertions(+), 5 deletions(-) create mode 100644 .github/workflows/run-ci-command.yml create mode 100644 .github/workflows/scripts/run_ci_command.py create mode 100644 .github/workflows/scripts/test_run_ci_command.py diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index 4124583d96d..a2f09eb8a45 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -80,9 +80,9 @@ jobs: '', '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in `#pr-reviews`, coordinate on features in `#feat-` channels, or join special interest groups in `#sig-` channels.', '', - 'PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.', + 'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.', '', - 'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.', + 'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.', '', 'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.', '', diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 143fc427a49..aa1f437ab6a 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -41,7 +41,7 @@ jobs: if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) { core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`); } else { - core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`); + core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found ${mergedCount}).`); } pre-commit: diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml new file mode 100644 index 00000000000..ff55c65dd75 --- /dev/null +++ b/.github/workflows/run-ci-command.yml @@ -0,0 +1,40 @@ +name: Run CI from PR comment + +on: + issue_comment: + types: [created] + +concurrency: + group: run-ci-comment-${{ github.event.issue.number }} + cancel-in-progress: false + +permissions: + contents: read + issues: write + pull-requests: read + +jobs: + run-ci-command: + if: >- + github.event.issue.pull_request && + (github.event.comment.body == '/ci run' || + github.event.comment.body == '/ci retry') + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: ${{ github.event.repository.default_branch }} + persist-credentials: false + - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + with: + python-version: "3.12" + - name: Authorize and run CI command + run: >- + uv run --no-project --python 3.12 + .github/workflows/scripts/run_ci_command.py + env: + BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} + BUILDKITE_ORGANIZATION: vllm + BUILDKITE_PIPELINE: ci + CI_TRUSTED_USERS: ${{ vars.CI_TRUSTED_USERS }} + GH_TOKEN: ${{ github.token }} diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py new file mode 100644 index 00000000000..7880fc5568a --- /dev/null +++ b/.github/workflows/scripts/run_ci_command.py @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from collections.abc import Mapping, Sequence +from typing import Any + +COMMAND_RUN_CI = "/ci run" +COMMAND_RETRY_FAILED = "/ci retry" +READY_LABELS = {"ready", "ready-run-all-tests"} +TRUSTED_PERMISSIONS = {"admin", "maintain", "write"} +ACTIVE_BUILD_STATES = { + "blocked", + "creating", + "scheduled", + "running", + "failing", + "canceling", + "waiting", + "waiting_failed", +} +RETRY_STATES = "failed,timed_out,expired" + + +class ApiError(RuntimeError): + def __init__(self, status: int | None, message: str) -> None: + super().__init__(message) + self.status = status + + +class HttpTransport: + def request( + self, + url: str, + *, + body: Mapping[str, Any] | None = None, + headers: Mapping[str, str] | None = None, + method: str = "GET", + ) -> Any: + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + url, + data=data, + headers=dict(headers or {}), + method=method, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + response_body = response.read().decode() + except urllib.error.HTTPError as error: + response_body = error.read().decode() + message = self._error_message(response_body, error.reason) + raise ApiError( + error.code, + f"API returned {error.code}: {message}", + ) from error + except urllib.error.URLError as error: + raise ApiError(None, f"API request failed: {error.reason}") from error + + if not response_body: + return None + try: + return json.loads(response_body) + except json.JSONDecodeError as error: + raise ApiError(None, "API returned a non-JSON response.") from error + + @staticmethod + def _error_message(response_body: str, fallback: str) -> str: + try: + parsed = json.loads(response_body) + except json.JSONDecodeError: + return fallback + return str(parsed.get("message", fallback)) + + +class GitHubClient: + def __init__( + self, + token: str, + repository: str, + transport: HttpTransport | None = None, + ) -> None: + if not token: + raise RuntimeError("GH_TOKEN is not set.") + self.owner, self.repo = repository.split("/", maxsplit=1) + self.transport = transport or HttpTransport() + self.headers = { + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": "vllm-ci-command", + "X-GitHub-Api-Version": "2022-11-28", + } + + def _request( + self, + path: str, + *, + body: Mapping[str, Any] | None = None, + method: str = "GET", + ) -> Any: + return self.transport.request( + f"https://api.github.com{path}", + body=body, + headers=self.headers, + method=method, + ) + + def _repo_path(self, suffix: str) -> str: + owner = urllib.parse.quote(self.owner, safe="") + repo = urllib.parse.quote(self.repo, safe="") + return f"/repos/{owner}/{repo}{suffix}" + + def _paginate(self, path: str) -> list[dict[str, Any]]: + results: list[dict[str, Any]] = [] + separator = "&" if "?" in path else "?" + for page in range(1, 101): + response = self._request(f"{path}{separator}per_page=100&page={page}") + if not isinstance(response, list): + raise ApiError(None, "GitHub API returned an invalid list response.") + results.extend(response) + if len(response) < 100: + return results + raise ApiError(None, "GitHub API pagination exceeded 10,000 results.") + + def get_pr(self, number: int) -> dict[str, Any]: + return self._request(self._repo_path(f"/pulls/{number}")) + + def get_permission(self, actor: str) -> str: + username = urllib.parse.quote(actor, safe="") + try: + response = self._request( + self._repo_path(f"/collaborators/{username}/permission") + ) + except ApiError as error: + if error.status == 404: + return "none" + raise + return str(response["permission"]) + + def get_review_decision(self, number: int) -> str | None: + query = """ + query($owner: String!, $repo: String!, $number: Int!) { + repository(owner: $owner, name: $repo) { + pullRequest(number: $number) { + reviewDecision + } + } + } + """ + response = self._request( + "/graphql", + body={ + "query": query, + "variables": { + "number": number, + "owner": self.owner, + "repo": self.repo, + }, + }, + method="POST", + ) + return response["data"]["repository"]["pullRequest"]["reviewDecision"] + + def list_reviews(self, number: int) -> list[dict[str, Any]]: + return self._paginate(self._repo_path(f"/pulls/{number}/reviews")) + + def list_reactions(self, comment_id: int) -> list[dict[str, Any]]: + return self._paginate( + self._repo_path(f"/issues/comments/{comment_id}/reactions") + ) + + def add_reaction(self, comment_id: int, content: str) -> None: + self._request( + self._repo_path(f"/issues/comments/{comment_id}/reactions"), + body={"content": content}, + method="POST", + ) + + def add_comment(self, issue_number: int, body: str) -> None: + self._request( + self._repo_path(f"/issues/{issue_number}/comments"), + body={"body": body}, + method="POST", + ) + + +class BuildkiteClient: + def __init__( + self, + token: str, + organization: str, + pipeline: str, + transport: HttpTransport | None = None, + ) -> None: + self.token = token + self.transport = transport or HttpTransport() + organization = urllib.parse.quote(organization, safe="") + pipeline = urllib.parse.quote(pipeline, safe="") + self.base_url = ( + "https://api.buildkite.com/v2/organizations/" + f"{organization}/pipelines/{pipeline}/builds" + ) + + def _request( + self, + *, + body: Mapping[str, Any] | None = None, + method: str = "GET", + path: str = "", + query: Sequence[tuple[str, str]] = (), + ) -> Any: + if not self.token: + raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.") + url = f"{self.base_url}{path}" + if query: + url = f"{url}?{urllib.parse.urlencode(query)}" + return self.transport.request( + url, + body=body, + headers={ + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json", + "User-Agent": "vllm-ci-command", + }, + method=method, + ) + + def list_builds( + self, + commit: str, + *, + metadata: tuple[str, str] | None = None, + ) -> list[dict[str, Any]]: + query = [ + ("commit", commit), + ("exclude_jobs", "true"), + ("exclude_pipeline", "true"), + ("per_page", "100"), + ] + if metadata: + key, value = metadata + query.append((f"meta_data[{key}]", value)) + response = self._request(query=query) + if not isinstance(response, list): + raise ApiError(None, "Buildkite API returned an invalid build list.") + return response + + def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]: + return self._request(body=body, method="POST") + + def retry_failed_jobs( + self, + build_number: int, + states: str, + ) -> dict[str, Any]: + number = urllib.parse.quote(str(build_number), safe="") + return self._request( + body={"states": states}, + method="PUT", + path=f"/{number}/retry_failed_jobs", + ) + + +def parse_command(body: str) -> str | None: + if body in {COMMAND_RUN_CI, COMMAND_RETRY_FAILED}: + return body + return None + + +def parse_trusted_users(value: str = "") -> set[str]: + return { + user.casefold() for item in value.split(",") for user in item.split() if user + } + + +def has_ready_label(pr: Mapping[str, Any]) -> bool: + return any(label["name"] in READY_LABELS for label in pr["labels"]) + + +def is_trusted_permission(permission: str) -> bool: + return permission in TRUSTED_PERMISSIONS + + +def authorize( + *, + actor: str, + permission: str, + pr: Mapping[str, Any], + trusted_approval: bool = False, + trusted_users: set[str] | None = None, +) -> tuple[bool, str]: + trusted_users = trusted_users or set() + if is_trusted_permission(permission): + return True, f"repository {permission} permission" + if actor.casefold() in trusted_users: + return True, "configured trusted contributor" + if actor.casefold() != pr["user"]["login"].casefold(): + return ( + False, + "Only reviewers with write access can run CI before it is " + "delegated to the PR author.", + ) + if pr["draft"]: + return False, "PR authors cannot run CI while the PR is a draft." + if has_ready_label(pr): + return True, "ready label" + if trusted_approval: + return True, "approval from a trusted reviewer" + return ( + False, + "A reviewer with write access must run `/ci run`, approve the PR, " + "or add the `ready` label first.", + ) + + +def has_trusted_approval( + github: GitHubClient, + number: int, + trusted_users: set[str], +) -> bool: + if github.get_review_decision(number) != "APPROVED": + return False + + latest_review_states: dict[str, tuple[str, str]] = {} + for review in github.list_reviews(number): + user = review.get("user") or {} + login = user.get("login") + state = review.get("state") + if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}: + latest_review_states[login.casefold()] = (login, state) + + for login, state in latest_review_states.values(): + if state != "APPROVED": + continue + if login.casefold() in trusted_users: + return True + if is_trusted_permission(github.get_permission(login)): + return True + return False + + +def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool: + pull_request = build.get("pull_request") + if isinstance(pull_request, Mapping): + build_pr_number = pull_request.get("id", pull_request.get("number")) + if build_pr_number is not None: + return str(build_pr_number) == str(pr_number) + metadata = build.get("meta_data") or {} + return str(metadata.get("github-pr-number")) == str(pr_number) + + +def is_active_build(build: Mapping[str, Any]) -> bool: + return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES + + +def select_latest_build( + builds: Sequence[dict[str, Any]], + pr_number: int, +) -> dict[str, Any] | None: + matching = [build for build in builds if is_build_for_pr(build, pr_number)] + return max(matching, key=lambda build: build.get("created_at", ""), default=None) + + +def create_build_payload( + *, + actor: str, + comment_id: int, + pr: Mapping[str, Any], +) -> dict[str, Any]: + return { + "commit": pr["head"]["sha"], + "branch": pr["head"]["ref"], + "message": f"PR #{pr['number']} {COMMAND_RUN_CI} by @{actor}", + "pull_request_id": pr["number"], + "pull_request_base_branch": pr["base"]["ref"], + "pull_request_repository": pr["head"]["repo"]["clone_url"], + "pull_request_labels": [label["name"] for label in pr["labels"]], + "env": { + "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id), + "VLLM_CI_TRIGGERED_BY": actor, + }, + "meta_data": { + "github-comment-id": str(comment_id), + "github-pr-number": str(pr["number"]), + "github-triggered-by": actor, + }, + } + + +def add_reaction_safely( + github: GitHubClient, + comment_id: int, + content: str, +) -> None: + try: + github.add_reaction(comment_id, content) + except Exception as error: + print(f"Could not add {content} reaction: {error}", file=sys.stderr) + + +def is_already_handled(github: GitHubClient, comment_id: int) -> bool: + return any( + reaction.get("content") in {"rocket", "-1"} + and (reaction.get("user") or {}).get("login") == "github-actions[bot]" + for reaction in github.list_reactions(comment_id) + ) + + +def handle_run_ci( + *, + actor: str, + buildkite: BuildkiteClient, + comment_id: int, + github: GitHubClient, + pr: Mapping[str, Any], +) -> str: + duplicate_builds = buildkite.list_builds( + pr["head"]["sha"], + metadata=("github-comment-id", str(comment_id)), + ) + duplicate = select_latest_build(duplicate_builds, pr["number"]) + if duplicate: + return f"CI was already requested by this comment: {duplicate['web_url']}" + + current_builds = buildkite.list_builds(pr["head"]["sha"]) + active_build = next( + ( + build + for build in current_builds + if is_build_for_pr(build, pr["number"]) and is_active_build(build) + ), + None, + ) + if active_build: + return f"CI is already running for this commit: {active_build['web_url']}" + + current_pr = github.get_pr(pr["number"]) + if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: + return ( + "The PR head changed while processing the command. Comment `/ci run` again." + ) + + build = buildkite.create_build( + create_build_payload( + actor=actor, + comment_id=comment_id, + pr=current_pr, + ) + ) + return ( + f"Triggered [Buildkite CI #{build['number']}]({build['web_url']}) " + f"for commit `{current_pr['head']['sha'][:12]}`." + ) + + +def handle_retry_failed( + *, + buildkite: BuildkiteClient, + pr: Mapping[str, Any], +) -> str: + builds = buildkite.list_builds(pr["head"]["sha"]) + build = select_latest_build(builds, pr["number"]) + if not build: + return "No CI build exists for the current PR commit. Use `/ci run` first." + if not build.get("finished_at") or is_active_build(build): + return f"CI is still running for this commit: {build['web_url']}" + + retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES) + if retried["retried_jobs_count"] == 0: + return ( + f"No failed, timed-out, or expired jobs need retrying: {build['web_url']}" + ) + return ( + f"Queued {retried['retried_jobs_count']} failed job(s) for retry in " + f"[Buildkite CI #{build['number']}]({build['web_url']})." + ) + + +def run( + event: Mapping[str, Any], + github: GitHubClient, + buildkite: BuildkiteClient, + trusted_users_value: str = "", +) -> None: + command = parse_command(event["comment"]["body"]) + if not command or "pull_request" not in event["issue"]: + return + + issue_number = event["issue"]["number"] + comment_id = event["comment"]["id"] + actor = event["comment"]["user"]["login"] + + if is_already_handled(github, comment_id): + print(f"Comment {comment_id} was already handled.") + return + add_reaction_safely(github, comment_id, "eyes") + + try: + pr = github.get_pr(issue_number) + permission = github.get_permission(actor) + if pr["state"] != "open": + github.add_comment(issue_number, "CI commands require an open PR.") + return + + trusted_users = parse_trusted_users(trusted_users_value) + should_check_approval = ( + not is_trusted_permission(permission) + and actor.casefold() not in trusted_users + and actor.casefold() == pr["user"]["login"].casefold() + and not pr["draft"] + and not has_ready_label(pr) + ) + trusted_approval = should_check_approval and has_trusted_approval( + github, + issue_number, + trusted_users, + ) + allowed, reason = authorize( + actor=actor, + permission=permission, + pr=pr, + trusted_approval=trusted_approval, + trusted_users=trusted_users, + ) + if not allowed: + add_reaction_safely(github, comment_id, "-1") + github.add_comment(issue_number, f"@{actor}, {reason}") + return + + print(f"Authorized @{actor}: {reason}") + if command == COMMAND_RUN_CI: + message = handle_run_ci( + actor=actor, + buildkite=buildkite, + comment_id=comment_id, + github=github, + pr=pr, + ) + else: + message = handle_retry_failed(buildkite=buildkite, pr=pr) + add_reaction_safely(github, comment_id, "rocket") + github.add_comment(issue_number, message) + except Exception: + add_reaction_safely(github, comment_id, "confused") + raise + + +def main() -> None: + event_path = os.environ["GITHUB_EVENT_PATH"] + with open(event_path, encoding="utf-8") as event_file: + event = json.load(event_file) + + if not parse_command(event["comment"]["body"]): + return + + github = GitHubClient( + os.environ.get("GH_TOKEN", ""), + os.environ["GITHUB_REPOSITORY"], + ) + buildkite = BuildkiteClient( + os.environ.get("BUILDKITE_API_TOKEN", ""), + os.environ.get("BUILDKITE_ORGANIZATION", "vllm"), + os.environ.get("BUILDKITE_PIPELINE", "ci"), + ) + run( + event, + github, + buildkite, + os.environ.get("CI_TRUSTED_USERS", ""), + ) + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py new file mode 100644 index 00000000000..e691438d6af --- /dev/null +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import unittest +from typing import Any + +from run_ci_command import ( + COMMAND_RETRY_FAILED, + COMMAND_RUN_CI, + RETRY_STATES, + BuildkiteClient, + authorize, + create_build_payload, + has_trusted_approval, + is_active_build, + is_build_for_pr, + parse_command, + parse_trusted_users, + run, + select_latest_build, +) + + +def make_pr(**overrides: Any) -> dict[str, Any]: + pr = { + "base": {"ref": "main"}, + "draft": False, + "head": { + "ref": "feature", + "repo": {"clone_url": "https://github.com/contributor/vllm.git"}, + "sha": "0123456789abcdef", + }, + "labels": [], + "number": 42, + "state": "open", + "user": {"login": "author"}, + } + pr.update(overrides) + return pr + + +def make_event(command: str, actor: str = "reviewer") -> dict[str, Any]: + return { + "comment": { + "body": command, + "id": 99, + "user": {"login": actor}, + }, + "issue": { + "number": 42, + "pull_request": {}, + }, + } + + +class FakeGitHub: + def __init__( + self, + *, + permission: str = "write", + permissions: dict[str, str] | None = None, + pr: dict[str, Any] | None = None, + review_decision: str = "REVIEW_REQUIRED", + reviews: list[dict[str, Any]] | None = None, + ) -> None: + self.comments: list[str] = [] + self.permission = permission + self.permissions = permissions or {} + self.pr = pr or make_pr() + self.reactions: list[str] = [] + self.review_decision = review_decision + self.reviews = reviews or [] + + def get_pr(self, number: int) -> dict[str, Any]: + return self.pr + + def get_permission(self, actor: str) -> str: + return self.permissions.get(actor, self.permission) + + def get_review_decision(self, number: int) -> str: + return self.review_decision + + def list_reviews(self, number: int) -> list[dict[str, Any]]: + return self.reviews + + def list_reactions(self, comment_id: int) -> list[dict[str, Any]]: + return [] + + def add_reaction(self, comment_id: int, content: str) -> None: + self.reactions.append(content) + + def add_comment(self, issue_number: int, body: str) -> None: + self.comments.append(body) + + +class FakeBuildkite: + def __init__( + self, + build_lists: list[list[dict[str, Any]]] | None = None, + ) -> None: + self.build_lists = build_lists or [] + self.created_builds: list[dict[str, Any]] = [] + self.list_calls: list[tuple[str, tuple[str, str] | None]] = [] + self.retry_calls: list[tuple[int, str]] = [] + + def list_builds( + self, + commit: str, + *, + metadata: tuple[str, str] | None = None, + ) -> list[dict[str, Any]]: + self.list_calls.append((commit, metadata)) + return self.build_lists.pop(0) + + def create_build(self, body: dict[str, Any]) -> dict[str, Any]: + self.created_builds.append(body) + return { + "number": 123, + "web_url": "https://buildkite.example/builds/123", + } + + def retry_failed_jobs( + self, + build_number: int, + states: str, + ) -> dict[str, Any]: + self.retry_calls.append((build_number, states)) + return {"retried_jobs_count": 3} + + +class FakeTransport: + def __init__(self, response: Any) -> None: + self.calls: list[dict[str, Any]] = [] + self.response = response + + def request(self, url: str, **kwargs: Any) -> Any: + self.calls.append({"url": url, **kwargs}) + return self.response + + +class RunCiCommandTest(unittest.TestCase): + def test_only_exact_ci_commands_are_accepted(self) -> None: + self.assertEqual(parse_command(COMMAND_RUN_CI), COMMAND_RUN_CI) + self.assertEqual( + parse_command(COMMAND_RETRY_FAILED), + COMMAND_RETRY_FAILED, + ) + self.assertIsNone(parse_command("/ci run please")) + self.assertIsNone(parse_command(" /ci run")) + + def test_write_access_authorizes_reviewers_and_authors(self) -> None: + allowed, _ = authorize( + actor="reviewer", + permission="write", + pr=make_pr(), + ) + self.assertTrue(allowed) + + def test_configured_trusted_contributors_can_run_ci(self) -> None: + trusted_users = parse_trusted_users("trusted-one, TRUSTED-TWO") + allowed, _ = authorize( + actor="trusted-two", + permission="read", + pr=make_pr(), + trusted_users=trusted_users, + ) + self.assertTrue(allowed) + + def test_authors_need_an_approval_or_ready_label(self) -> None: + pending, _ = authorize( + actor="author", + permission="read", + pr=make_pr(), + ) + approved, _ = authorize( + actor="author", + permission="read", + pr=make_pr(), + trusted_approval=True, + ) + ready, _ = authorize( + actor="author", + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + self.assertFalse(pending) + self.assertTrue(approved) + self.assertTrue(ready) + + def test_non_author_contributors_without_write_are_denied(self) -> None: + allowed, _ = authorize( + actor="contributor", + permission="read", + pr=make_pr(), + trusted_approval=True, + ) + self.assertFalse(allowed) + + def test_authors_cannot_use_ready_state_on_draft_prs(self) -> None: + allowed, _ = authorize( + actor="author", + permission="read", + pr=make_pr(draft=True, labels=[{"name": "ready"}]), + trusted_approval=True, + ) + self.assertFalse(allowed) + + def test_only_trusted_reviewers_can_delegate_through_approval(self) -> None: + approved_review = { + "state": "APPROVED", + "user": {"login": "reviewer"}, + } + trusted = FakeGitHub( + permission="read", + permissions={"reviewer": "write"}, + review_decision="APPROVED", + reviews=[approved_review], + ) + untrusted = FakeGitHub( + permission="read", + review_decision="APPROVED", + reviews=[approved_review], + ) + self.assertTrue(has_trusted_approval(trusted, 42, set())) + self.assertFalse(has_trusted_approval(untrusted, 42, set())) + + def test_build_matching_is_scoped_to_the_pr(self) -> None: + self.assertTrue(is_build_for_pr({"pull_request": {"id": 42}}, 42)) + self.assertFalse(is_build_for_pr({"pull_request": {"id": 43}}, 42)) + self.assertTrue( + is_build_for_pr( + {"meta_data": {"github-pr-number": "42"}}, + 42, + ) + ) + + def test_latest_build_selection_ignores_other_prs(self) -> None: + latest = select_latest_build( + [ + { + "created_at": "2026-07-28T02:00:00Z", + "number": 3, + "pull_request": {"id": 43}, + }, + { + "created_at": "2026-07-28T01:00:00Z", + "number": 2, + "pull_request": {"id": 42}, + }, + { + "created_at": "2026-07-28T00:00:00Z", + "number": 1, + "pull_request": {"id": 42}, + }, + ], + 42, + ) + self.assertEqual(latest["number"], 2) + + def test_active_build_states_prevent_duplicate_runs(self) -> None: + self.assertTrue(is_active_build({"state": "scheduled"})) + self.assertTrue(is_active_build({"state": "running"})) + self.assertTrue(is_active_build({"state": "waiting"})) + self.assertTrue(is_active_build({"blocked": True, "state": "passed"})) + self.assertFalse(is_active_build({"state": "failed"})) + + def test_build_payload_preserves_pr_context(self) -> None: + payload = create_build_payload( + actor="reviewer", + comment_id=99, + pr=make_pr(labels=[{"name": "ready"}, {"name": "v1"}]), + ) + self.assertEqual( + payload, + { + "commit": "0123456789abcdef", + "branch": "feature", + "message": "PR #42 /ci run by @reviewer", + "pull_request_id": 42, + "pull_request_base_branch": "main", + "pull_request_repository": ("https://github.com/contributor/vllm.git"), + "pull_request_labels": ["ready", "v1"], + "env": { + "VLLM_CI_GITHUB_COMMENT_ID": "99", + "VLLM_CI_TRIGGERED_BY": "reviewer", + }, + "meta_data": { + "github-comment-id": "99", + "github-pr-number": "42", + "github-triggered-by": "reviewer", + }, + }, + ) + + def test_ci_run_dispatches_build_with_current_pr_metadata(self) -> None: + github = FakeGitHub() + buildkite = FakeBuildkite([[], []]) + run(make_event(COMMAND_RUN_CI), github, buildkite) + + self.assertEqual(len(buildkite.created_builds), 1) + self.assertEqual( + buildkite.created_builds[0]["message"], + "PR #42 /ci run by @reviewer", + ) + self.assertEqual(github.reactions, ["eyes", "rocket"]) + self.assertIn("Buildkite CI #123", github.comments[0]) + + def test_unapproved_authors_are_denied_without_buildkite(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(), + review_decision="REVIEW_REQUIRED", + ) + buildkite = FakeBuildkite() + run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) + + self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.reactions, ["eyes", "-1"]) + self.assertIn("approve the PR", github.comments[0]) + + def test_ci_retry_uses_latest_current_sha_build(self) -> None: + github = FakeGitHub( + permission="read", + pr=make_pr(labels=[{"name": "ready"}]), + ) + buildkite = FakeBuildkite( + [ + [ + { + "created_at": "2026-07-28T01:00:00Z", + "finished_at": "2026-07-28T02:00:00Z", + "number": 123, + "pull_request": {"id": 42}, + "state": "failed", + "web_url": "https://buildkite.example/builds/123", + } + ] + ] + ) + run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) + + self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)]) + self.assertIn("Queued 3 failed job", github.comments[0]) + + def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None: + transport = FakeTransport({"retried_jobs_count": 2}) + client = BuildkiteClient( + "secret", + "vllm", + "ci", + transport=transport, + ) + client.retry_failed_jobs(123, RETRY_STATES) + + call = transport.calls[0] + self.assertEqual(call["method"], "PUT") + self.assertTrue(call["url"].endswith("/123/retry_failed_jobs")) + self.assertEqual(call["body"], {"states": RETRY_STATES}) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 34dc385db78..89acc6b7f5a 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -301,8 +301,10 @@ review process: isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion. - Note that not all CI checks will be executed due to limited computational - resources. The reviewer will add `ready` label to the PR when the PR is - ready to merge or a full CI run is needed. + resources. Reviewers with write access and configured trusted contributors + can comment `/ci run` when CI signals are needed before a PR is ready. After + the PR is approved or has the `ready` label, the PR author can use `/ci run` + or `/ci retry`. New commits do not start CI automatically. ### Pull Request Limits and Escalation From 54ab69b14eeafd5c8dcf818cfcc0eeb26d2ebddb Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 28 Jul 2026 19:11:04 -0700 Subject: [PATCH 54/67] [CI] Allow comment-triggered builds past pipeline filters (#50197) Signed-off-by: khluu --- .github/workflows/scripts/run_ci_command.py | 1 + .github/workflows/scripts/test_run_ci_command.py | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 7880fc5568a..5bf196bb53f 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -381,6 +381,7 @@ def create_build_payload( "pull_request_base_branch": pr["base"]["ref"], "pull_request_repository": pr["head"]["repo"]["clone_url"], "pull_request_labels": [label["name"] for label in pr["labels"]], + "ignore_pipeline_branch_filters": True, "env": { "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id), "VLLM_CI_TRIGGERED_BY": actor, diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index e691438d6af..edbfff3f5f0 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -280,6 +280,7 @@ class RunCiCommandTest(unittest.TestCase): "pull_request_base_branch": "main", "pull_request_repository": ("https://github.com/contributor/vllm.git"), "pull_request_labels": ["ready", "v1"], + "ignore_pipeline_branch_filters": True, "env": { "VLLM_CI_GITHUB_COMMENT_ID": "99", "VLLM_CI_TRIGGERED_BY": "reviewer", From 17a74b745b459eb72ea9eea2945bc83b1c17e0a0 Mon Sep 17 00:00:00 2001 From: Krishna Teja Chitty-Venkata <44275589+krishnateja95@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:27:01 -0400 Subject: [PATCH 55/67] [Model] Add Inkling compressed-tensors dynamic FP8 support (#48876) Signed-off-by: mgoin Co-authored-by: mgoin Co-authored-by: OpenAI Codex --- .../models/inkling/test_moe_weight_layout.py | 26 +++++++++++++++++++ vllm/models/inkling/nvidia/moe.py | 3 +++ 2 files changed, 29 insertions(+) diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py index a7307e167a1..e15a573b845 100644 --- a/tests/models/inkling/test_moe_weight_layout.py +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -187,6 +187,32 @@ def test_moe_loads_compressed_tensors_global_scale( assert loaded == [f"experts.routed_experts.{projection}_{scale_kind}_global_scale"] +@pytest.mark.parametrize(("projection", "checkpoint_rows"), [("w13", 8), ("w2", 4)]) +def test_moe_loads_channelwise_scale_for_tp( + projection: str, checkpoint_rows: int +) -> None: + param = torch.nn.Parameter(torch.empty(2, 4, 1)) + experts = SimpleNamespace( + **{f"{projection}_weight_scale": param}, + moe_config=SimpleNamespace(moe_parallel_config=SimpleNamespace(tp_rank=1)), + ) + layer = SimpleNamespace( + experts=SimpleNamespace(routed_experts=experts), + _local_expert_slots=lambda: {0: 0, 2: 1}, + ) + checkpoint_scale = torch.arange(3 * checkpoint_rows).reshape(3, checkpoint_rows, 1) + + loaded = moe.InklingMoE.load_expert_weight( + layer, f"experts.{projection}_weight_scale", checkpoint_scale + ) + + expected = checkpoint_scale[[0, 2]] + if projection == "w13": + expected = expected[:, 4:].reshape(2, 2, 2, 1).transpose(1, 2).flatten(1, 2) + torch.testing.assert_close(param, expected.float()) + assert loaded == [f"experts.routed_experts.{projection}_weight_scale"] + + def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None: monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 32f7489a4c4..9ff3acaa14d 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/vllm/models/inkling/nvidia/moe.py @@ -589,6 +589,9 @@ class InklingMoE(nn.Module): param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to( param.device ) + elif key == "w2_weight_scale" and weight.shape[-1] == 1: + # Per-output-channel scales are replicated across TP ranks. + param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype) elif key.startswith("w13"): # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the # fused param layout is [w1(gate); w3(up)]. The TP-local rows form From 7398a30d79758559704d4f28173b3594b2ec1345 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 28 Jul 2026 21:29:00 -0500 Subject: [PATCH 56/67] [ROCm][CI] Stabilize ngram and suffix correctness test (#50190) Signed-off-by: Andreas Karatzas --- tests/v1/e2e/spec_decode/test_spec_decode.py | 27 ++++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 17c48721f80..482cbfb82ed 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -23,7 +23,7 @@ from vllm import LLM, SamplingParams from vllm.assets.base import VLLM_S3_BUCKET_URL from vllm.assets.image import VLM_IMAGES_DIR from vllm.benchmarks.datasets import InstructCoderDataset -from vllm.config import VllmConfig, replace +from vllm.config import CompilationConfig, VllmConfig, replace from vllm.distributed import cleanup_dist_env_and_memory from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform @@ -160,6 +160,12 @@ def reset_torch_dynamo(): torch._dynamo.reset() +@pytest.fixture +def disable_vllm_compile_cache_on_rocm(request: pytest.FixtureRequest) -> None: + if current_platform.is_rocm(): + request.getfixturevalue("disable_vllm_compile_cache") + + @pytest.mark.parametrize( "speculative_config", [ @@ -175,21 +181,26 @@ def reset_torch_dynamo(): }, ], ) +@pytest.mark.usefixtures("disable_vllm_compile_cache_on_rocm") @single_gpu_only @large_gpu_mark(min_gb=20) def test_ngram_and_suffix_correctness( speculative_config: dict, model_name: str, + vllm_runner, ): - spec_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + # Keep LLM defaults; VllmRunner only provides lifecycle cleanup here. + trust_remote_code=False, + enable_chunked_prefill=None, speculative_config=speculative_config, max_model_len=4096, - ) - evaluate_llm_for_gsm8k(spec_llm) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + # Preserve LLM's default compilation/cudagraph configuration. Without + # this, VllmRunner injects its reduced test-only capture sizes. + compilation_config=CompilationConfig(), + ) as runner: + evaluate_llm_for_gsm8k(runner.llm) @pytest.mark.parametrize("async_scheduling", [True], ids=["async"]) From 30c2718eaafc230927f0d8ac47439d040b46a7c3 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 28 Jul 2026 22:34:20 -0400 Subject: [PATCH 57/67] [CompressedTensors] FP4 Qutlass Integration (#43229) Signed-off-by: Kyle Sayers Signed-off-by: Brian Dellabetta Signed-off-by: Brian Dellabetta Co-authored-by: Brian Dellabetta Co-authored-by: Brian Dellabetta Co-authored-by: Dipika Sikka --- tests/quantization/test_compressed_tensors.py | 5 + vllm/_custom_ops.py | 37 +++++- vllm/model_executor/layers/linear.py | 1 + .../compressed_tensors/transform/linear.py | 6 +- .../compressed_tensors/transform/module.py | 36 +++--- .../transform/schemes/linear_qutlass_nvfp4.py | 107 +++++++++++++++++- .../layers/quantization/qutlass_utils.py | 11 +- 7 files changed, 169 insertions(+), 34 deletions(-) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 626717cd4a3..8c509888700 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -473,6 +473,11 @@ def test_compressed_tensors_w4a8_fp8(vllm_runner, args): "Flat is better than nested.\nSparse is better than dense.", 150.0, ), + ( + "nm-testing/Llama-3.2-1B-Instruct-quipv16-nvfp4", + "Flat is better than nested.\nSparse is better than dense.", + 150.0, + ), ], ) def test_compressed_tensors_transforms_perplexity( diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index b1d2f1344d6..9f592386c15 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3994,9 +3994,40 @@ def fusedQuantizeNv( padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device ) - return torch.ops._qutlass_C.fusedQuantizeNvAbsMax( - a, b, xh_e2m1, xh_e4m3, global_scale - ) + safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale) + return xh_e2m1, xh_e4m3 + + +@torch.library.custom_op( + "vllm::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3") +) +def safeFusedQuantizeNv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, +) -> None: + """ + Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place + rather than returning them, to prevent torch 2.12+ errors that outputs of custom + operators may not alias any inputs to the custom operator. + """ + torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale) + return + + +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"): + + @register_fake("vllm::safeFusedQuantizeNv") + def _fake_fused_quantize_nv( + a: torch.Tensor, + b: torch.Tensor, + xh_e2m1: torch.Tensor, + xh_e4m3: torch.Tensor, + global_scale: torch.Tensor, + ) -> None: + return def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor: diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index df105f06634..e4662148ff7 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -50,6 +50,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [ "UnquantizedLinearMethod", "CompressedTensorsLinearMethod", "CompressedTensorsLinearTransformMethod", + "QutlassNvFP4LinearMethod", "AutoAWQMarlinLinearMethod", "AutoAWQLinearMethod", "AutoGPTQLinearMethod", diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py index bd1964e667d..0cde020627d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py @@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.transform.module from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501 TransformTuple, ) +from vllm.platforms import current_platform class CompressedTensorsLinearTransformMethod(LinearMethodBase): @@ -48,11 +49,12 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase): assert input_tfms or output_tfms - if is_qutlass_fp4_scheme(quant_scheme, input_tfms): + if is_qutlass_fp4_scheme( + quant_scheme, input_tfms + ) and current_platform.has_device_capability(100): return QutlassNvFP4LinearMethod(quant_method, input_tfms, output_tfms) # hadacore or dense gemm is selected by Transform module - return cls(quant_method, input_tfms, output_tfms) def __init__( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py index f5589c8c07f..75505934938 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py @@ -32,7 +32,7 @@ class HadamardTransform(torch.nn.Module): transforms: dict[int, TransformTuple] # info parsed from transforms config weight: SharedWeightParameter # container for shared tensors - scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0)) + scaled_data_ptrs: set[int] = set() def __init__( self, @@ -44,7 +44,6 @@ class HadamardTransform(torch.nn.Module): ): super().__init__() self.transforms = transforms - self.scales = {} if get_tensor_model_parallel_world_size() > 1: raise NotImplementedError( @@ -64,26 +63,26 @@ class HadamardTransform(torch.nn.Module): ) data_key = self._get_data_key(scheme, weight_size) + # load up in model's default precision, rather than using scheme.precision self.weight.add_partition( part_index, data_key, size=(weight_size, weight_size), - dtype=scheme.precision, ) # validate that shared tensors and schemes are correct self._validate_input_transforms() def process_weights_after_loading(self): - for part_id in self.weight.partitions: - data = self.weight.partitions[part_id].data - + for part_id, partition in self.weight.partitions.items(): # required by torch.compile self.weight.process_weights_after_loading() - # precompute scale as a runtime multiply, not division - # do not fold into weight in order to utilize FWHT - self.scales[part_id] = 1 / math.sqrt(data.size(0)) + # Merge normalization scale directly into weight, must be done only once + data_ptr = partition.data.data_ptr() + if data_ptr not in HadamardTransform.scaled_data_ptrs: + partition.data.div_(math.sqrt(partition.data.size(0))) + HadamardTransform.scaled_data_ptrs.add(data_ptr) # FUTURE: avoid runtime transpose by processing weights # prior to apply @@ -111,26 +110,19 @@ class HadamardTransform(torch.nn.Module): weight = ( weight if self.transforms[part_id].args.inverse else weight.T ) # linear := x(W.T) - scale = self.scales[part_id] if self.transforms[part_id].scheme.head_dim is not None: value = value.unflatten(-1, (-1, weight.size(0))) - value = ( - dispatch_unquantized_gemm()( - self, value.to(weight.dtype), weight, None - ).to(value.dtype) - * scale - ) + value = dispatch_unquantized_gemm()( + self, value.to(weight.dtype), weight, None + ).to(value.dtype) value = value.flatten(-2, -1) return value - return ( - dispatch_unquantized_gemm()( - self, value.to(weight.dtype), weight, None - ).to(value.dtype) - * scale - ) + return dispatch_unquantized_gemm()( + self, value.to(weight.dtype), weight, None + ).to(value.dtype) def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable: return (id(scheme), weight_size) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py index f0bb47a728a..b8b771e3c7a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py @@ -2,7 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from torch.nn.parameter import Parameter +from vllm._custom_ops import fusedQuantizeNv +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + NvFp4LinearKernel, +) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 CompressedTensorsScheme, CompressedTensorsW4A4Fp4, @@ -11,18 +17,32 @@ from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear CompressedTensorsLinearTransformMethod, TransformTuple, ) +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + slice_nvfp4_output, +) +from vllm.utils.flashinfer import ( + flashinfer_scaled_fp4_mm, +) __all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"] +NVFP4_MAX = 6.0 + +# QUTLASS supports transform block sizes (16, 32, 64, 128) for NVFP4 +# https://github.com/IST-DASLab/qutlass/blob/v0.2.0/qutlass/csrc/bindings.cpp#L413-L414 def is_qutlass_fp4_scheme( quant_scheme: CompressedTensorsScheme | None, input_tfms: dict[int, TransformTuple], ) -> bool: return ( - isinstance(quant_scheme, (CompressedTensorsW4A4Fp4,)) - and len(input_tfms) == 1 - and input_tfms[0].scheme.head_dim == quant_scheme.group_size + isinstance(quant_scheme, CompressedTensorsW4A4Fp4) + and len(input_tfms) >= 1 + and all( + input_tfm.scheme.head_dim in (16, 32, 64, 128) + for input_tfm in input_tfms.values() + ) ) @@ -50,15 +70,90 @@ class QutlassNvFP4LinearMethod(CompressedTensorsLinearTransformMethod): ) assert self.input_transform is not None - assert len(self.input_transform.weight) == 1 - assert self.input_transform.weight[0].size(0) == layer.scheme.group_size + assert len(self.input_transform.weight.partitions) >= 1 return ret + @staticmethod + def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str: + """ + Given a kernel, find the string that is needed to be passed into + `flashinfer_scaled_fp4_mm`, using + vllm.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth + """ + kernel_type = type(kernel) + for key, kernels in _LINEAR_BACKEND_KERNEL_MAP.items(): + if not key.startswith("flashinfer_") or kernel_type not in kernels: + continue + backend = key.removeprefix("flashinfer_") + # flashinfer GEMM backend uses "cute-dsl", not "cutedsl" + return backend.replace("cutedsl", "cute-dsl") + raise ValueError( + f"QutlassNvFP4 transform requires a FlashInfer kernel, " + f"got {kernel_type.__name__}" + ) + + def process_weights_after_loading(self, layer): + super().process_weights_after_loading(layer) + + assert self.input_transform is not None + layer.hadamard_matrix = self.input_transform.weight.partitions[0].data + + # fusedQuantizeNv stores raw absmax as block scales (sf = absmax), + # while CT weights use sf = absmax * SFScaleVal / 6.0. The GEMM + # computes alpha * sum(fp4_a * sf_a * fp4_w * sf_w), so alpha must + # compensate: alpha = weight_global_scale / 6.0 + layer.fused_alpha = Parameter( + layer.weight_global_scale / NVFP4_MAX, requires_grad=False + ) + + layer.fused_global_scale = Parameter( + torch.tensor( + [NVFP4_MAX], + dtype=torch.float32, + device=layer.weight_global_scale.device, + ), + requires_grad=False, + ) + + layer.flashinfer_gemm_backend = self._get_flashinfer_gemm_backend( + layer.scheme.kernel + ) + def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - raise NotImplementedError() + assert bias is None + output_size = layer.output_size_per_partition + output_shape = [*x.shape[:-1], output_size] + + x_flat = x.contiguous().flatten(end_dim=-2) + + x_fp4, x_scales = fusedQuantizeNv( + x_flat, layer.hadamard_matrix, layer.fused_global_scale + ) + + x_scales_blocked = to_blocked(x_scales, backend="triton").view(x_scales.shape) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_scales_blocked, + layer.weight_scale, + layer.fused_alpha, + x.dtype, + backend=layer.flashinfer_gemm_backend, + ) + + out = slice_nvfp4_output(out, output_size) + + if self.output_transform is not None: + for part_id, (start, length) in enumerate(self.partition_ranges): + out[:, start : start + length] = self.output_transform( + out[:, start : start + length].clone(), part_id=part_id + ) + + return out.view(*output_shape) diff --git a/vllm/model_executor/layers/quantization/qutlass_utils.py b/vllm/model_executor/layers/quantization/qutlass_utils.py index 315ecd0c009..86b0548307a 100644 --- a/vllm/model_executor/layers/quantization/qutlass_utils.py +++ b/vllm/model_executor/layers/quantization/qutlass_utils.py @@ -84,6 +84,7 @@ def triton_scale_swizzle( ) +@torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=()) def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: """ Rearranges an E8M0 tensor scale from row-major format to @@ -142,6 +143,14 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor: return out +@triton_mx_block_rearrange.register_fake +def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor: + rows, cols = scale_tensor.shape + padded_rows = cdiv(rows, 128) * 128 + padded_cols = cdiv(cols, 4) * 4 + return scale_tensor.new_empty((padded_rows, padded_cols)) + + def to_blocked( input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton" ) -> torch.Tensor: @@ -157,7 +166,7 @@ def to_blocked( backend: "torch" (PyTorch path) or "triton" (Triton kernel) Returns: - Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4)) + Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4)) """ if backend == "triton": return triton_mx_block_rearrange(input_matrix).flatten() From 32a423ac0aad67f94f93e97f73338f484b55faec Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Wed, 29 Jul 2026 05:46:38 +0300 Subject: [PATCH 58/67] Integrate CuTeDSL MoE for ReLU2 NVFP4 (#49580) Co-authored-by: OpenAI Codex Co-authored-by: Michael Goin --- .../moe/test_flashinfer_cutedsl_nvfp4_moe.py | 230 ++++++++++++++++++ .../experts/flashinfer_cutedsl_moe.py | 8 +- .../quantization/utils/flashinfer_fp4_moe.py | 16 +- 3 files changed, 245 insertions(+), 9 deletions(-) create mode 100644 tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py new file mode 100644 index 00000000000..a7a7c5251bc --- /dev/null +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for FlashInfer CuTeDSL NVFP4 MoE.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.quantization.nvfp4_utils import ( + FLOAT4_E2M1_MAX, + FLOAT8_E4M3_MAX, + break_fp4_bytes, +) +from tests.kernels.utils import torch_moe +from vllm import _custom_ops as ops +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import fused_topk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( + FlashInferCuteDSLExperts, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4 +from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import set_random_seed + +if not has_flashinfer_cutedsl_moe_nvfp4() or not ( + current_platform.is_device_capability_family(100) +): + pytest.skip( + "Requires FlashInfer CuTeDSL NVFP4 MoE on SM100", + allow_module_level=True, + ) + + +def _quantize_nvfp4_linear( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + weights_q = [] + scales = [] + global_scales = [] + for expert_weight in weight: + global_scale = ( + FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max() + ).to(torch.float32) + weight_q, scale = ops.scaled_fp4_quant( + expert_weight, + global_scale, + is_sf_swizzled_layout=False, + ) + weights_q.append(weight_q) + scales.append(scale) + global_scales.append(global_scale) + return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales) + + +def _dequantize_nvfp4_linear( + tensor_fp4: torch.Tensor, + tensor_sf: torch.Tensor, + global_scale: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + assert tensor_fp4.dtype == torch.uint8 + m, packed_k = tensor_fp4.shape + k = packed_k * 2 + tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32) + tensor_f32 = tensor_f32.reshape(m, k // 16, 16) + tensor_sf = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32) + tensor_sf = tensor_sf[:, : k // 16] / global_scale + return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype) + + +@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.inference_mode() +def test_flashinfer_cutedsl_fp4_moe_relu2_no_mul( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + + w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) + w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + + activation = MoEActivation.RELU2_NO_MUL + fake_layer = SimpleNamespace(activation=activation) + a1_scale = torch.ones(1, device="cuda", dtype=torch.float32) + a2_scale = torch.ones(1, device="cuda", dtype=torch.float32) + ( + w1_cutedsl, + w1_scale_cutedsl, + w1_alpha, + a1_scale, + w2_cutedsl, + w2_scale_cutedsl, + w2_alpha, + a2_scale, + ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( + layer=fake_layer, + w13=w1_q, + w13_scale=w1_scale, + w13_scale_2=(1.0 / w1_global_scale), + a13_scale=a1_scale, + w2=w2_q, + w2_scale=w2_scale, + w2_scale_2=(1.0 / w2_global_scale), + a2_scale=a2_scale, + ) + quant_config = nvfp4_moe_quant_config( + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + a1_gscale=(1.0 / a1_scale), + a2_gscale=(1.0 / a2_scale), + w1_scale=w1_scale_cutedsl, + w2_scale=w2_scale_cutedsl, + is_scale_swizzled=False, + ) + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=dtype, + routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), + ) + + cutedsl_experts = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + FlashInferCuteDSLExperts( + moe_config=moe_config, + quant_config=quant_config, + ), + ) + + cutedsl_output = cutedsl_experts.apply( + hidden_states=hidden_states, + w1=w1_cutedsl, + w2=w2_cutedsl, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=e, + expert_map=None, + apply_router_weight_on_input=False, + ) + + a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32) + a_q, a_scale = ops.scaled_fp4_quant( + hidden_states, + a_global_scale, + is_sf_swizzled_layout=False, + ) + a_in_dtype = _dequantize_nvfp4_linear( + a_q, + a_scale, + a_global_scale, + dtype=dtype, + ) + + w1_d = torch.empty((e, n, k), device="cuda", dtype=dtype) + w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) + for idx in range(e): + w1_d[idx] = _dequantize_nvfp4_linear( + w1_q[idx], + w1_scale[idx], + w1_global_scale[idx], + dtype=dtype, + ) + w2_d[idx] = _dequantize_nvfp4_linear( + w2_q[idx], + w2_scale[idx], + w2_global_scale[idx], + dtype=dtype, + ) + + torch_output = torch_moe( + a_in_dtype, + w1_d, + w2_d, + score, + topk, + activation=activation, + ) + torch.testing.assert_close( + torch_output, + cutedsl_output, + atol=2e-1, + rtol=2e-1, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index b512d51c135..823ae43b5ef 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -13,6 +13,9 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, @@ -76,7 +79,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( @@ -90,7 +93,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU + return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) @staticmethod def _supports_parallel_config( @@ -163,4 +166,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): num_local_experts=self.local_num_experts, local_expert_offset=self.local_expert_offset, moe_output=output, + activation_type=activation_to_flashinfer_int(activation), ) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index bab3dee649b..1b513e9c0ce 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -103,7 +103,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend. Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper, - and interleaves w13 gate/linear rows. + and interleaves w13 gate/linear rows for gated activations. Non-gated + activations use a single w13 projection and keep its row order unchanged. """ from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout @@ -112,13 +113,14 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) - half = w13.shape[1] // 2 - w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) - w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) + if layer.activation.is_gated: + half = w13.shape[1] // 2 + w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) + w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) - # Interleave up/gate rows for w13 weights and scales. - w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) - w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) + # Interleave up/gate rows for w13 weights and scales. + w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) + w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) # Convert w13 scale factors: linear → swizzled → MMA layout. w13_scale = swizzle_blockscale(w13_scale) From f37f03db4af69e4969242767734db3d9175055f7 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 20:50:30 -0700 Subject: [PATCH 59/67] [KV Connector] Support NIXL P/D for hybrid MLA+SSM models (#49762) Signed-off-by: Nick Hill Co-authored-by: Jared Wen --- .../unit/test_nixl_connector_hma.py | 196 ++++++++++++++++++ .../unit/test_nixl_desc_geometry.py | 21 ++ .../unit/test_nixl_push_connector.py | 4 + .../kv_connector/v1/nixl/base_worker.py | 30 ++- .../kv_connector/v1/nixl/push_worker.py | 68 +++--- 5 files changed, 281 insertions(+), 38 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 1f7a62d2c9a..0124f6516ef 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1380,3 +1380,199 @@ def test_logical_to_kernel_block_ids_with_remote_ratio( assert list(result) == expected_kernel_block_ids, ( f"Expected {expected_kernel_block_ids}, got {result}" ) + + +# ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ───────────────────── + + +def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4): + """KimiLinear-shaped config: one MLA group and two KDA (GDN-typed + MambaSpec) groups whose layers share the same HMA tensors, with a + mamba-aligned unified page and an MLA kernel block smaller than the + logical block.""" + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + # 12-token logical blocks over a 4-token MLA kernel block. + mla_spec = MLAAttentionSpec( + block_size=12, num_kv_heads=1, head_size=6, dtype=torch.float16 + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=12, + # GDN-decomposable conv (Q|K|V = 2|2|4 cols x 3 rows) + fp32 temporal. + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + assert kda_spec.page_size_bytes == unified_page + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + +@pytest.mark.cpu_test +def test_register_kv_caches_hybrid_mla_dual_purpose_regions(): + """Hybrid MLA+KDA registration: HMA tensors shared by both layer types + must be flagged as MLA regions even when a KDA layer registers them + first, expose TP-independent kernel-granularity block lens, and build + FA + mamba descriptors for every region.""" + from unittest.mock import MagicMock + + from vllm.config import set_current_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + kv_cache_config = _make_hybrid_mla_kv_cache_config() + unified_page = kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes + vllm_config = create_vllm_config(block_size=12) + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared tensors would not be + # deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [4] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + with ( + patch.object(bw, "NixlWrapper"), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "test-engine", kv_cache_config) + worker.use_mla = True # opt-125m test config is not MLA; force the flag + worker.nixl_wrapper.get_agent_metadata.return_value = b"fake-agent-metadata" + + tensors = [torch.zeros(4 * unified_page, dtype=torch.uint8) for _ in range(2)] + # KDA layer first per tensor: exercises the dual-purpose flag merge. + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + + # 12-token logical blocks over the 4-token MLA kernel block. + assert worker._physical_blocks_per_logical_kv_block == 3 + assert worker.block_size == 4 and worker.num_blocks == 12 + # Both shared tensors are dual-purpose: their FA view is MLA even though + # a KDA layer registered them first. + assert worker._region_is_mla == [True, True] + assert worker.num_regions == 2 and worker.num_descs == 24 + # Kernel-granularity block lens; TP-independent for MLA hybrids. + assert worker.block_len_per_layer == [unified_page // 3] * 2 + # Split handles must replicate every FA descriptor (MLA isn't head-sharded). + assert worker._fa_desc_replicated(worker.num_descs) == [True] * 24 + # FA descs: 2 regions x 12 kernel blocks, page stride = kernel page. + # Mamba descs: 2 regions x (3 conv sub-projections + 1 ssm) x 4 blocks. + assert worker.src_blocks_data.shape == (24 + 32, 3) + fa_descs = worker.src_blocks_data[:24] + assert fa_descs[1][0] - fa_descs[0][0] == unified_page // 3 + assert all(size == unified_page // 3 for size in fa_descs[:, 1]) + + +@pytest.mark.cpu_test +def test_push_write_hybrid_mla_replicates_attention(): + """Hybrid MLA+SSM push with P_TP < D_TP: attention blocks must be + written to every covered D rank (replicated MLA latent) while SSM state + is written per-rank through the split handles.""" + import threading + from collections import defaultdict + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + ) + from vllm.v1.kv_cache_interface import MambaSpec, MLAAttentionSpec + + worker = object.__new__(NixlPushConnectorWorker) + worker.shutdown = lambda: None # skeleton worker: silence __del__ + worker.use_mla = True + worker._has_mamba = True + worker._group_spec_types = (MLAAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.transfer_topo.tp_ratio.return_value = -2 + remote_info = MagicMock() + remote_info.remote_physical_blocks_per_logical = 1 + remote_info.remote_block_size = 4 + worker.transfer_topo.get_engine_info.return_value = remote_info + + engine_id = "remote-engine" + # Read-oriented mapping collapses the replicated attention group to one + # source rank; the SSM state is sharded across both covered D ranks. + worker.tp_mappings = { + engine_id: TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + } + worker.dst_xfer_side_handles = {engine_id: {0: 100, 1: 101}} + worker.src_xfer_handles_by_tp_ratio = {(-2, 4): [200, 201]} + worker.src_xfer_handles_by_block_size = {4: 300} + worker._sending_transfers = defaultdict(list) + worker._sending_transfers_lock = threading.Lock() + worker.kv_cache_config = _make_hybrid_mla_kv_cache_config() + worker._xfer_blocks = MagicMock(return_value=1) + + meta = MagicMock() + meta.remote.engine_id = engine_id + meta.remote.block_ids = [[7, 8], [3]] + meta.local_physical_block_ids = [[1, 2], [5]] + + worker._xfer_blocks_for_req("req-1", meta) + + calls = worker._xfer_blocks.call_args_list + assert len(calls) == 2 + for call, rank, local_handle, remote_handle in zip( + calls, (0, 1), (200, 201), (100, 101) + ): + spec = call.kwargs["read_spec"] + assert spec.remote_rank == rank + # Attention group replicated to every rank, SSM by membership. + assert spec.local_block_ids == [[1, 2], [5]] + assert spec.remote_block_ids == [[7, 8], [3]] + assert call.kwargs["local_xfer_side_handle"] == local_handle + assert call.kwargs["remote_xfer_side_handle"] == remote_handle diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index e2fc41a0229..6ac9ce030d0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -617,3 +617,24 @@ def test_mla_hybrid_large_ppl_geometry(num_tokens): num_tokens=num_tokens, tp_size=8, ) + + +@pytest.mark.cpu_test +def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid(): + """The MLA per-token page is TP-independent, so kernel block lengths + differing by anything other than the block-size ratio must fail the + handshake loudly rather than transfer at mismatched geometry.""" + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + # Equal kernel block sizes (ratio 1), but a half-sized per-token page. + meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer] + with pytest.raises((AssertionError, RuntimeError)): + worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index d670e5563b6..adf0063701c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -43,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) +from vllm.v1.kv_cache_interface import FullAttentionSpec from vllm.v1.outputs import KVConnectorOutput from .utils import make_nixl_push_scheduler @@ -337,6 +338,9 @@ class _StubWriterWorker(NixlPushConnectorWorker): w.engine_id = "test-decode-engine" w._remote_agents = {} w._physical_blocks_per_logical_kv_block = 1 + # Single non-hybrid attention group, matching the stub block id lists. + w._has_mamba = False + w._group_spec_types = (FullAttentionSpec,) # Track _do_start_push_kv invocations. calls: list[tuple[str, Any, dict[str, Any]]] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 480ac8937df..12b9f0d937f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1093,7 +1093,7 @@ class NixlBaseConnectorWorker: caches_data = [] # With hybrid allocator, layers can share a kv cache tensor - seen_base_addresses = [] + seen_base_addresses: list[int] = [] # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead @@ -1139,11 +1139,19 @@ class NixlBaseConnectorWorker: curr_tensor_size_bytes = num_blocks * physical_page_size base_addr = cache.data_ptr() + is_mla_region = isinstance( + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) + ) if base_addr in seen_base_addresses: # NOTE (NickLucche) HMA employs memory pooling to share tensors # across groups. This results in skipping all tensors but the ones # pointed to by group0. Also, generally we will have more blocks # per tensor but fewer regions. + # A shared tensor may back both SSM and attention layers (e.g. + # KDA+MLA in KimiLinear); the region's FA view is MLA whichever + # layer registered it first. + idx = seen_base_addresses.index(base_addr) + self._region_is_mla[idx] |= is_mla_region logger.debug("Skipping %s because it's already seen", layer_name) continue logger.debug( @@ -1157,9 +1165,6 @@ class NixlBaseConnectorWorker: ) else: self.block_len_per_layer.append(physical_page_size) - is_mla_region = isinstance( - layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) - ) self._region_is_mla.append(is_mla_region) if not is_mla_region: @@ -1801,7 +1806,22 @@ class NixlBaseConnectorWorker: # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. - if not self._has_mamba: + if self._has_mamba and self.use_mla: + # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are + # kernel-granularity views of the mamba-unified page. The MLA + # per-token page is TP-independent, so the block lengths must + # match up to the kernel block size ratio even under + # heterogeneous TP (remote kernel blocks may be smaller). + # SSM geometry is validated via ssm_sizes/conv offsets instead. + assert self.block_len_per_layer == [ + block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens + ], ( + "Hybrid MLA kernel-granularity block lengths must match " + f"between P and D (block_size_ratio={block_size_ratio}): " + f"local={self.block_len_per_layer}, " + f"remote={nixl_agent_meta.block_lens}." + ) + elif not self._has_mamba: assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 6bc48963c58..88977e34033 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -52,7 +52,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( ReqMeta, TransferHandle, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + ReadSpec, + _is_attention_spec, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id from vllm.logger import init_logger @@ -505,41 +508,37 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): local_block_ids = meta.local_physical_block_ids num_groups = len(local_block_ids) - if self.use_mla and tp_ratio < 0: - # MLA latent is replicated across D's TP ranks: the tp-mapping - # collapses to one rank (fine for reads), but push must WRITE every - # D rank or the rest decode stale KV; only the dst differs per rank. + # MLA latent is replicated across D's TP ranks: the tp-mapping + # collapses it to one rank (fine for reads), but push must WRITE every + # D rank or the rest decode stale KV. For hybrid MLA+SSM the sharded + # SSM state already targets every covered D rank, so only the + # attention groups need widening; pure MLA writes to all handshaked + # ranks (only the dst differs per rank). + replicate_attn = self.use_mla and tp_ratio < 0 + if replicate_attn and not self._has_mamba: assert len(plan.all_source_ranks) == 1 - mla_local_ids = [list(ids) for ids in local_block_ids] - mla_remote_ids = [list(ids) for ids in remote_block_ids] - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=mla_local_ids, - remote_block_ids=mla_remote_ids, - ) - for rank in self.dst_xfer_side_handles[engine_id] - ] + write_ranks = sorted(self.dst_xfer_side_handles[engine_id]) else: - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=[ - list(local_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - remote_block_ids=[ - list(remote_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - ) - for rank in plan.all_source_ranks + write_ranks = list(plan.all_source_ranks) + + def group_ids(block_ids: BlockIds, rank: int) -> BlockIds: + return [ + list(block_ids[g]) + if (replicate_attn and _is_attention_spec(self._group_spec_types[g])) + or rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) ] + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=group_ids(local_block_ids, rank), + remote_block_ids=group_ids(remote_block_ids, rank), + ) + for rank in write_ranks + ] + handles: list[int] = [] for i, spec in enumerate(read_specs): remote_block_size = remote_info.remote_block_size @@ -551,7 +550,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): remote_block_size, req_id, ) - if tp_ratio < 0 and not self.use_mla: + if tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1): + # Multiple targets: write each rank its chunk of local memory. + # Hybrid MLA+SSM also lands here: its split handles replicate + # the attention descriptors and chunk only the SSM state. split_key = (tp_ratio, remote_block_size) local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: From 58f96593971b6287cdc8867024a6dfb9985c545c Mon Sep 17 00:00:00 2001 From: Zach Zhu Date: Wed, 29 Jul 2026 12:20:40 +0800 Subject: [PATCH 60/67] [Frontend][Core] Standardize request error handling with VLLMError hierarchy (#49665) Signed-off-by: Zach Zhu --- tests/engine/test_short_mm_context.py | 6 +- tests/entrypoints/llm/test_chat.py | 3 +- .../entrypoints/llm/test_prompt_validation.py | 5 +- .../multimodal/llm/test_mm_embeds_only.py | 3 +- .../openai/chat_completion/test_chat.py | 7 +- .../openai/chat_completion/test_chat_error.py | 10 +- .../chat_completion/test_logprob_token_ids.py | 8 +- .../test_thinking_token_budget_validation.py | 6 +- .../completion/test_completion_error.py | 8 +- .../completion/test_prompt_validation.py | 3 +- .../openai/responses/test_sampling_params.py | 3 +- .../test_http_status_metrics.py | 168 ++++++++++-------- .../entrypoints/unit_tests/test_chat_utils.py | 5 +- .../multimodal/processing/test_gemma4.py | 3 +- .../processing/test_gemma4_unified.py | 3 +- .../media/test_unprocessable_entity_error.py | 6 +- tests/multimodal/test_processing.py | 7 +- .../test_chat_utils_prompt_embeds.py | 5 +- tests/renderers/test_completions.py | 7 +- tests/samplers/test_non_finite_params.py | 2 +- tests/test_envs.py | 5 +- tests/test_pooling_params.py | 28 +-- tests/test_sampling_params.py | 3 +- ...est_chat_completion_request_validations.py | 11 +- .../test_responses_request_validations.py | 12 +- tests/v1/e2e/general/test_context_length.py | 2 +- tests/v1/e2e/general/test_min_tokens.py | 5 +- tests/v1/e2e/general/test_streaming_input.py | 13 +- tests/v1/engine/test_async_llm.py | 7 +- tests/v1/logits_processors/utils.py | 3 +- tests/v1/sample/test_logprobs.py | 2 +- tests/v1/sample/test_sampling_params_e2e.py | 7 +- tests/v1/structured_output/test_validation.py | 5 +- vllm/entrypoints/openai/api_server.py | 36 ++-- vllm/entrypoints/openai/engine/protocol.py | 4 +- .../entrypoints/serve/utils/error_response.py | 23 ++- vllm/entrypoints/serve/utils/server_utils.py | 12 +- vllm/exceptions.py | 24 ++- vllm/inputs/engine.py | 6 +- vllm/pooling_params.py | 13 +- vllm/sampling_params.py | 72 ++++---- vllm/v1/engine/async_llm.py | 11 +- vllm/v1/engine/exceptions.py | 7 +- vllm/v1/engine/input_processor.py | 27 +-- vllm/v1/sample/logits_processor/__init__.py | 8 +- vllm/v1/sample/logits_processor/interface.py | 4 +- 46 files changed, 370 insertions(+), 248 deletions(-) diff --git a/tests/engine/test_short_mm_context.py b/tests/engine/test_short_mm_context.py index 23489c21333..940709c8e53 100644 --- a/tests/engine/test_short_mm_context.py +++ b/tests/engine/test_short_mm_context.py @@ -3,6 +3,8 @@ import pytest +from vllm.exceptions import VLLMValidationError + from ..conftest import IMAGE_ASSETS HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( @@ -19,7 +21,9 @@ models = ["llava-hf/llava-1.5-7b-hf"] def test_context_length_too_short(vllm_runner, image_assets, model): images = [asset.pil_image for asset in image_assets] - with pytest.raises(ValueError, match="longer than the maximum model length"): + with pytest.raises( + VLLMValidationError, match="longer than the maximum model length" + ): vllm_model = vllm_runner( model, # LLaVA has a feature size of 576 diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index 61cdbd3eee2..cbc57b80da4 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -6,6 +6,7 @@ import pytest from vllm import LLM from vllm.distributed import cleanup_dist_env_and_memory +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams @@ -157,7 +158,7 @@ def test_chat_batch_failure_cleanup(llm_for_failure_test): batch_2 = [valid_msg, valid_msg] sampling_params = SamplingParams(temperature=0, max_tokens=10) - with pytest.raises(ValueError, match="maximum context length is"): + with pytest.raises(VLLMValidationError, match="maximum context length is"): llm.chat(batch_1, sampling_params=sampling_params) assert llm.llm_engine.get_num_unfinished_requests() == 0 diff --git a/tests/entrypoints/llm/test_prompt_validation.py b/tests/entrypoints/llm/test_prompt_validation.py index c17486d962f..8dd55c6b10e 100644 --- a/tests/entrypoints/llm/test_prompt_validation.py +++ b/tests/entrypoints/llm/test_prompt_validation.py @@ -5,17 +5,18 @@ import pytest import torch from vllm import LLM +from vllm.exceptions import VLLMValidationError def test_empty_prompt(): llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(ValueError, match="decoder prompt cannot be empty"): + with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"): llm.generate([""]) def test_out_of_vocab_token(): llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(ValueError, match="out of vocabulary"): + with pytest.raises(VLLMValidationError, match="out of vocabulary"): llm.generate({"prompt_token_ids": [999999]}) diff --git a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py index 57bec9c1188..0ea18071246 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py +++ b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py @@ -6,6 +6,7 @@ import pytest from tests.entrypoints.multimodal.conftest import managed_llm from vllm import LLM, SamplingParams from vllm.assets.image import ImageAsset +from vllm.exceptions import VLLMValidationError MODEL = "llava-hf/llava-1.5-7b-hf" PROMPT = "USER: \nDescribe this image briefly.\nASSISTANT:" @@ -42,7 +43,7 @@ def test_generate_with_embedding(llm: LLM): def test_raw_image_rejected(llm: LLM): """Raw image input is still rejected when limit=0.""" raw_image = ImageAsset("stop_sign").pil_image - with pytest.raises(ValueError, match=r"At most 0 image\(s\)"): + with pytest.raises(VLLMValidationError, match=r"At most 0 image\(s\)"): llm.generate( {"prompt": PROMPT, "multi_modal_data": {"image": raw_image}}, sampling_params=SamplingParams(max_tokens=16), diff --git a/tests/entrypoints/openai/chat_completion/test_chat.py b/tests/entrypoints/openai/chat_completion/test_chat.py index 32c72f1ef93..5541b605d99 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_chat.py @@ -18,6 +18,7 @@ from tests.utils import RemoteOpenAIServer from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams # any model with a chat template should work here @@ -1074,7 +1075,7 @@ def test_chat_completion_request_n_parameter_exceeds_default_limit( max_tokens=10, ) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): request.to_sampling_params( max_tokens=10, default_sampling_params={}, @@ -1136,7 +1137,7 @@ def test_chat_completion_request_n_parameter_custom_limit( max_tokens=10, ) - with pytest.raises(ValueError, match="n must be at most 128"): + with pytest.raises(VLLMValidationError, match="n must be at most 128"): request_over.to_sampling_params( max_tokens=10, default_sampling_params={}, @@ -1160,7 +1161,7 @@ def test_chat_completion_request_n_parameter_massive_value( max_tokens=1, ) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): request.to_sampling_params( max_tokens=1, default_sampling_params={}, diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 4b42f522e81..9e805254a2f 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -6,7 +6,6 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest -from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -18,6 +17,7 @@ from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.scale_out.render.serving import ServingRender +from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -479,7 +479,7 @@ def test_json_schema_response_format_missing_schema(): def test_structural_tag_response_format_invalid(format_value): """Malformed structural tags should be rejected during request validation.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): ChatCompletionRequest( @@ -493,7 +493,7 @@ def test_structural_tag_response_format_invalid(format_value): def test_batch_structural_tag_response_format_invalid(format_value): """Batch chat should reject malformed structural tags at request parsing.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): BatchChatCompletionRequest( @@ -507,7 +507,7 @@ def test_batch_structural_tag_response_format_invalid(format_value): def test_structured_outputs_structural_tag_invalid(structural_tag): """Malformed direct structured_outputs structural tags should be rejected.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid structured_outputs structural_tag", ): ChatCompletionRequest( @@ -521,7 +521,7 @@ def test_structured_outputs_structural_tag_invalid(structural_tag): def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" - with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"): + with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): ChatCompletionRequest( model=MODEL_NAME, messages=[{"role": "user", "content": "hello"}], diff --git a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py index aa04d787ccc..1eea6db48a7 100644 --- a/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py +++ b/tests/entrypoints/openai/chat_completion/test_logprob_token_ids.py @@ -14,11 +14,11 @@ digit-token vocab id). import math import pytest -from pydantic import ValidationError from tests.utils import RemoteOpenAIServer from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.exceptions import VLLMValidationError MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" @@ -87,7 +87,7 @@ def test_completion_request_decouples_top_k_from_explicit_token_ids(): def test_completion_rejects_explicit_token_ids_without_generated_tokens(): - with pytest.raises(ValidationError, match="no output tokens are generated"): + with pytest.raises(VLLMValidationError, match="no output tokens are generated"): CompletionRequest( model=MODEL_NAME, prompt="Hello", @@ -99,7 +99,7 @@ def test_completion_rejects_explicit_token_ids_without_generated_tokens(): def test_requests_reject_explicit_token_ids_with_beam_search(): - with pytest.raises(ValidationError, match="not supported with beam search"): + with pytest.raises(VLLMValidationError, match="not supported with beam search"): ChatCompletionRequest( model=MODEL_NAME, messages=[{"role": "user", "content": "Hello"}], @@ -108,7 +108,7 @@ def test_requests_reject_explicit_token_ids_with_beam_search(): use_beam_search=True, ) - with pytest.raises(ValidationError, match="not supported with beam search"): + with pytest.raises(VLLMValidationError, match="not supported with beam search"): CompletionRequest( model=MODEL_NAME, prompt="Hello", diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py index e66205b7df2..1b2b76dc093 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py @@ -2,15 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from pydantic import ValidationError from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.exceptions import VLLMValidationError @pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5]) def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value): - with pytest.raises(ValidationError, match="thinking_token_budget"): + with pytest.raises(VLLMValidationError, match="thinking_token_budget"): ChatCompletionRequest.model_validate( { "model": "qwen", @@ -44,7 +44,7 @@ def test_chat_completion_request_accepts_minus_one_as_unlimited(): @pytest.mark.parametrize("raw_value", [0.6, 3.14, -2]) def test_completion_request_rejects_invalid_thinking_token_budget(raw_value): - with pytest.raises(ValidationError, match="thinking_token_budget"): + with pytest.raises(VLLMValidationError, match="thinking_token_budget"): CompletionRequest.model_validate( { "model": "qwen", diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 818cad738d9..2f3db9f697d 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -6,7 +6,6 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.completion.protocol import CompletionRequest @@ -18,6 +17,7 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.scale_out.render.serving import ServingRender +from vllm.exceptions import VLLMValidationError from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -430,7 +430,7 @@ def test_json_schema_response_format_missing_schema(): def test_structural_tag_response_format_invalid(format_value): """Malformed structural tags should be rejected during request validation.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid response_format structural_tag", ): CompletionRequest( @@ -445,7 +445,7 @@ def test_structural_tag_response_format_invalid(format_value): def test_structured_outputs_structural_tag_invalid(structural_tag): """Malformed direct structured_outputs structural tags should be rejected.""" with pytest.raises( - ValidationError, + VLLMValidationError, match="Invalid structured_outputs structural_tag", ): CompletionRequest( @@ -616,7 +616,7 @@ class TestCompletionPromptListLimit: def test_non_numeric_logprobs_rejected(field_name): """A non-numeric logprobs value must be a clean 400 validation error, not a TypeError from the mode='before' comparison (which surfaces as HTTP 500).""" - with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"): + with pytest.raises(VLLMValidationError, match=f"`{field_name}` must be an integer"): CompletionRequest( model=MODEL_NAME, prompt="Test prompt", diff --git a/tests/entrypoints/openai/completion/test_prompt_validation.py b/tests/entrypoints/openai/completion/test_prompt_validation.py index 87c6b6e1668..6c40037e07c 100644 --- a/tests/entrypoints/openai/completion/test_prompt_validation.py +++ b/tests/entrypoints/openai/completion/test_prompt_validation.py @@ -13,6 +13,7 @@ import torch from tests.utils import RemoteOpenAIServer from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.renderers.embed_utils import safe_load_prompt_embeds @@ -111,5 +112,5 @@ def test_disable_prompt_embeds(dtype: torch.dtype, seq_len: int, hidden_size: in buffer.seek(0) encoded_tensor = pybase64.b64encode(buffer.getvalue()) - with pytest.raises(ValueError, match="--enable-prompt-embeds"): + with pytest.raises(VLLMValidationError, match="--enable-prompt-embeds"): safe_load_prompt_embeds(model_config, encoded_tensor) diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index 5a68e3a9c0d..6ede3f1f7f2 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -14,6 +14,7 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ResponseTextConfig, ) +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import StructuredOutputsParams @@ -163,7 +164,7 @@ class TestResponsesRequestSamplingParams: text=text_config, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(VLLMValidationError) as exc_info: request.to_sampling_params(default_max_tokens=1000) assert "Cannot specify both structured_outputs and text.format" in str( diff --git a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py index 0f96bf161d2..060b05bbaca 100644 --- a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py @@ -8,54 +8,95 @@ PrometheusInstrumentatorMiddleware before being caught by ServerErrorMiddleware. """ from argparse import Namespace -from http import HTTPStatus import httpx import pytest -from fastapi import FastAPI, HTTPException, Request -from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse +from fastapi import HTTPException from prometheus_client import CollectorRegistry -from prometheus_fastapi_instrumentator import Instrumentator -from vllm.entrypoints.serve.utils.server_utils import exception_handler -from vllm.exceptions import VLLMNotFoundError, VLLMValidationError +from vllm.entrypoints.openai.api_server import build_app +from vllm.exceptions import ( + VLLMNotFoundError, + VLLMServerError, + VLLMValidationError, +) -@pytest.fixture +@pytest.fixture(scope="module") +def should_do_global_cleanup_after_test() -> bool: + # This suite never initializes distributed/accelerator state. + return False + + +def _build_args() -> Namespace: + """Minimal args for ``build_app``; avoids ``make_arg_parser`` device probing.""" + return Namespace( + disable_fastapi_docs=True, + enable_offline_docs=False, + root_path=None, + allowed_origins=["*"], + allow_credentials=False, + allowed_methods=["*"], + allowed_headers=["*"], + api_key=None, + enable_request_id_headers=False, + enable_fault_tolerance=False, + middleware=[], + log_error_stack=False, + ) + + +@pytest.fixture(scope="module") def registry(): - """Create a fresh Prometheus registry for each test.""" + """Shared Prometheus registry for the module-scoped app.""" return CollectorRegistry() -@pytest.fixture +@pytest.fixture(scope="module") def app(registry): - """Create a minimal FastAPI app that mirrors vLLM's exception handler - and Prometheus middleware setup.""" + """Build the real vLLM FastAPI app once and attach probe routes that raise. - app = FastAPI() + Patch the name used by ``attach_router`` (imported into the instrumentator + metrics module), not ``vllm.v1.metrics.prometheus`` alone — that binding is + captured at import time. + """ + import vllm.entrypoints.serve.instrumentator.metrics as metrics_mod - # Mock app state that exception_handler needs - app.state.args = Namespace(log_error_stack=False) + original = metrics_mod.get_prometheus_registry + metrics_mod.get_prometheus_registry = lambda: registry + try: + app = build_app(_build_args(), supported_tasks=()) + finally: + metrics_mod.get_prometheus_registry = original - # Register exception handlers exactly as vLLM does in build_app() - app.exception_handler(HTTPException)(_http_exception_handler) - app.exception_handler(RequestValidationError)(_validation_exception_handler) - app.exception_handler(ValueError)(exception_handler) - app.exception_handler(TypeError)(exception_handler) - app.exception_handler(OverflowError)(exception_handler) - app.exception_handler(NotImplementedError)(exception_handler) - app.exception_handler(VLLMValidationError)(exception_handler) - app.exception_handler(VLLMNotFoundError)(exception_handler) - app.exception_handler(Exception)(exception_handler) + @app.get("/raise_http_exception_400") + async def raise_http_exception_400(): + raise HTTPException(status_code=400, detail="bad request") - # Instrument with Prometheus (same as vLLM's attach_router) - Instrumentator( - excluded_handlers=["/metrics"], - registry=registry, - ).add().instrument(app) + @app.get("/raise_http_exception_404") + async def raise_http_exception_404(): + raise HTTPException(status_code=404, detail="not found") + + @app.get("/raise_request_validation_error") + async def raise_request_validation_error(n: int): + # Invalid ``n`` triggers FastAPI's RequestValidationError. + return {"n": n} + + @app.get("/raise_vllm_validation_error") + async def raise_vllm_validation_error(): + raise VLLMValidationError("bad parameter", parameter="temperature") + + @app.get("/raise_vllm_not_found_error") + async def raise_vllm_not_found_error(): + raise VLLMNotFoundError("model not found") + + @app.get("/raise_vllm_server_error") + async def raise_vllm_server_error(): + # Bare VLLMServerError goes through vllm_error_handler → 500. + # EngineGenerateError / EngineDeadError are not used here: they call + # terminate_if_errored and need engine/server state. + raise VLLMServerError("internal server failure") - # Test routes that raise different exception types @app.get("/raise_value_error") async def raise_value_error(): raise ValueError("invalid input value") @@ -72,22 +113,6 @@ def app(registry): async def raise_not_implemented_error(): raise NotImplementedError("feature not supported") - @app.get("/raise_vllm_validation_error") - async def raise_vllm_validation_error(): - raise VLLMValidationError("bad parameter", parameter="temperature") - - @app.get("/raise_vllm_not_found_error") - async def raise_vllm_not_found_error(): - raise VLLMNotFoundError("model not found") - - @app.get("/raise_http_exception_400") - async def raise_http_exception_400(): - raise HTTPException(status_code=400, detail="bad request") - - @app.get("/raise_http_exception_404") - async def raise_http_exception_404(): - raise HTTPException(status_code=404, detail="not found") - @app.get("/raise_runtime_error") async def raise_runtime_error(): raise RuntimeError("unexpected server error") @@ -99,14 +124,6 @@ def app(registry): return app -async def _http_exception_handler(req: Request, exc: HTTPException): - return JSONResponse({"error": exc.detail}, status_code=exc.status_code) - - -async def _validation_exception_handler(req: Request, exc: RequestValidationError): - return JSONResponse({"error": str(exc)}, status_code=HTTPStatus.BAD_REQUEST) - - def _get_http_requests_total(registry, method: str, handler: str): """Extract the http_requests_total metric values grouped by status. @@ -128,31 +145,31 @@ def _get_http_requests_total(registry, method: str, handler: str): @pytest.mark.asyncio @pytest.mark.parametrize( - "endpoint,expected_status_group,expected_http_code", + "endpoint,expected_status_group,expected_http_code,request_kwargs", [ - # These should record as 4xx in Prometheus - ("/raise_value_error", "4xx", 400), - ("/raise_type_error", "4xx", 400), - ("/raise_overflow_error", "4xx", 400), - ("/raise_vllm_validation_error", "4xx", 400), - ("/raise_vllm_not_found_error", "4xx", 404), - ("/raise_http_exception_400", "4xx", 400), - ("/raise_http_exception_404", "4xx", 404), - # NotImplementedError returns 501 which is still 5xx group - ("/raise_not_implemented_error", "5xx", 501), - # These should record as 5xx in Prometheus (genuine server errors) - ("/raise_runtime_error", "5xx", 500), - # Successful requests should record as 2xx - ("/success", "2xx", 200), + ("/raise_http_exception_400", "4xx", 400, {}), + ("/raise_http_exception_404", "4xx", 404, {}), + ("/raise_request_validation_error", "4xx", 400, {"params": {"n": "x"}}), + ("/raise_vllm_validation_error", "4xx", 400, {}), + ("/raise_vllm_not_found_error", "4xx", 404, {}), + ("/raise_vllm_server_error", "5xx", 500, {}), + ("/raise_value_error", "4xx", 400, {}), + ("/raise_type_error", "4xx", 400, {}), + ("/raise_overflow_error", "4xx", 400, {}), + ("/raise_not_implemented_error", "5xx", 501, {}), + ("/raise_runtime_error", "5xx", 500, {}), + ("/success", "2xx", 200, {}), ], ids=[ + "HTTPException(400)->4xx", + "HTTPException(404)->4xx", + "RequestValidationError->4xx", + "VLLMValidationError->4xx", + "VLLMNotFoundError->4xx", + "VLLMServerError->5xx", "ValueError->4xx", "TypeError->4xx", "OverflowError->4xx", - "VLLMValidationError->4xx", - "VLLMNotFoundError->4xx", - "HTTPException(400)->4xx", - "HTTPException(404)->4xx", "NotImplementedError->5xx", "RuntimeError->5xx", "success->2xx", @@ -164,6 +181,7 @@ async def test_http_requests_total_records_correct_status( endpoint, expected_status_group, expected_http_code, + request_kwargs, ): """Verify that http_requests_total records the correct status group. @@ -177,7 +195,7 @@ async def test_http_requests_total_records_correct_status( async with httpx.AsyncClient( transport=transport, base_url="http://testserver" ) as client: - response = await client.get(endpoint) + response = await client.get(endpoint, **request_kwargs) # Verify the HTTP response code returned to the client is correct assert response.status_code == expected_http_code, ( diff --git a/tests/entrypoints/unit_tests/test_chat_utils.py b/tests/entrypoints/unit_tests/test_chat_utils.py index 82b262321d6..b5a35c2cb25 100644 --- a/tests/entrypoints/unit_tests/test_chat_utils.py +++ b/tests/entrypoints/unit_tests/test_chat_utils.py @@ -21,6 +21,7 @@ from vllm.entrypoints.chat_utils import ( parse_chat_messages, parse_chat_messages_async, ) +from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from vllm.multimodal.utils import ( encode_audio_url, @@ -1504,7 +1505,7 @@ def test_parse_chat_messages_rejects_too_many_images_in_one_message( "ignore", message="coroutine 'async_get_and_parse_image' was never awaited", ) - with pytest.raises(ValueError, match="At most"): + with pytest.raises(VLLMValidationError, match="At most"): parse_chat_messages( [ { @@ -1540,7 +1541,7 @@ def test_parse_chat_messages_rejects_too_many_images_across_messages( "ignore", message="coroutine 'async_get_and_parse_image' was never awaited", ) - with pytest.raises(ValueError, match="At most"): + with pytest.raises(VLLMValidationError, match="At most"): parse_chat_messages( [ { diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py index a355501fdd8..f30afe47dde 100644 --- a/tests/models/multimodal/processing/test_gemma4.py +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -7,6 +7,7 @@ import pytest import torch from PIL import Image as PILImage +from vllm.exceptions import VLLMValidationError from vllm.model_executor.models.gemma4_mm import ( Gemma4ForConditionalGeneration, Gemma4ImagePixelInputs, @@ -222,7 +223,7 @@ def test_limit_mm_per_prompt( mm_data = {"image": images} # Expect ValueError when exceeding limit - with pytest.raises(ValueError, match="At most 1 image"): + with pytest.raises(VLLMValidationError, match="At most 1 image"): processor( prompt, mm_items=processor.info.parse_mm_data(mm_data), diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 473ba729b85..67a81ddb7b4 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -7,6 +7,7 @@ import pytest import torch from PIL import Image as PILImage +from vllm.exceptions import VLLMValidationError from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -197,7 +198,7 @@ def test_limit_mm_per_prompt( mm_data = {"image": images} - with pytest.raises(ValueError, match="At most 1 image"): + with pytest.raises(VLLMValidationError, match="At most 1 image"): processor( prompt, mm_items=processor.info.parse_mm_data(mm_data), diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 8be70383b8a..7cad4295575 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -14,7 +14,7 @@ import aiohttp import pytest from vllm.entrypoints.serve.utils.error_response import create_error_response -from vllm.exceptions import VLLMUnprocessableEntityError +from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector @@ -35,9 +35,9 @@ class TestVLLMUnprocessableEntityError: assert "parameter=image_url" in str(exc) assert "value=https://example.com/image.jpg" in str(exc) - def test_is_value_error_subclass(self): + def test_is_client_error_subclass(self): exc = VLLMUnprocessableEntityError("Test") - assert isinstance(exc, ValueError) + assert isinstance(exc, VLLMClientError) class TestMediaConnectorErrorHandling: diff --git a/tests/multimodal/test_processing.py b/tests/multimodal/test_processing.py index 66acdbe62ff..2153cd2c259 100644 --- a/tests/multimodal/test_processing.py +++ b/tests/multimodal/test_processing.py @@ -8,6 +8,7 @@ import numpy as np import pytest from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.processing.context import InputProcessingContext from vllm.multimodal.processing.processor import ( @@ -931,7 +932,11 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid): else: mm_data = {"image": [image] * num_images} - exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most") + exc_ctx = ( + nullcontext() + if is_valid + else pytest.raises(VLLMValidationError, match="At most") + ) with exc_ctx: processor( diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index 2238c41f498..3f08194b9be 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -26,6 +26,7 @@ from vllm.entrypoints.chat_utils import ( parse_chat_messages, parse_chat_messages_async, ) +from vllm.exceptions import VLLMValidationError from vllm.renderers.hf import ( _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, _build_mixed_prompt_embeds, @@ -264,7 +265,7 @@ def test_parse_chat_messages_requires_flag(): "content": [{"type": "prompt_embeds", "data": b64}], } ] - with pytest.raises(ValueError, match=_ENABLE_PROMPT_EMBEDS_ERROR): + with pytest.raises(VLLMValidationError, match=_ENABLE_PROMPT_EMBEDS_ERROR): parse_chat_messages( messages, mc, @@ -283,7 +284,7 @@ def test_parse_chat_messages_rejects_missing_data(): "content": [{"type": "prompt_embeds"}], # no `data` } ] - with pytest.raises(ValueError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): + with pytest.raises(VLLMValidationError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): parse_chat_messages( messages, mc, diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index d184eb8621c..1849eeac7dc 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -11,6 +11,7 @@ import pytest import torch from vllm.config import ModelConfig +from vllm.exceptions import VLLMValidationError from vllm.inputs import SingletonPrompt from vllm.renderers import TokenizeParams from vllm.renderers.hf import HfRenderer @@ -286,7 +287,7 @@ class TestRenderPrompt: ) with pytest.raises( - ValueError, + VLLMValidationError, match="maximum context length is", ): renderer.tokenize_prompts( @@ -307,7 +308,7 @@ class TestRenderPrompt: ) with pytest.raises( - ValueError, + VLLMValidationError, match="maximum context length is", ): renderer.tokenize_prompts( @@ -328,7 +329,7 @@ class TestRenderPrompt: ) with pytest.raises( - ValueError, + VLLMValidationError, match="maximum context length is", ): renderer.tokenize_prompts( diff --git a/tests/samplers/test_non_finite_params.py b/tests/samplers/test_non_finite_params.py index 57fe90f314c..f982953d608 100644 --- a/tests/samplers/test_non_finite_params.py +++ b/tests/samplers/test_non_finite_params.py @@ -42,7 +42,7 @@ class TestNonFiniteRepetitionPenalty: ids=["nan", "inf", "-inf", "math.nan", "math.inf"], ) def test_non_finite_repetition_penalty_rejected(self, value: float): - with pytest.raises(ValueError, match="repetition_penalty"): + with pytest.raises(VLLMValidationError, match="repetition_penalty"): SamplingParams(repetition_penalty=value) def test_finite_repetition_penalty_accepted(self): diff --git a/tests/test_envs.py b/tests/test_envs.py index 5917c28fab2..5e0363e33a1 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -15,6 +15,7 @@ from vllm.envs import ( env_with_choices, environment_variables, ) +from vllm.exceptions import VLLMValidationError def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch): @@ -547,7 +548,7 @@ class TestVllmMaxNSequences: max_n = envs.VLLM_MAX_N_SEQUENCES SamplingParams(n=max_n) - with pytest.raises(ValueError, match="n must be at most"): + with pytest.raises(VLLMValidationError, match="n must be at most"): SamplingParams(n=max_n + 1) def test_sampling_params_respects_custom_limit( @@ -563,5 +564,5 @@ class TestVllmMaxNSequences: SamplingParams(n=128) - with pytest.raises(ValueError, match="n must be at most 128"): + with pytest.raises(VLLMValidationError, match="n must be at most 128"): SamplingParams(n=129) diff --git a/tests/test_pooling_params.py b/tests/test_pooling_params.py index 17d04078b4e..f34270d0da5 100644 --- a/tests/test_pooling_params.py +++ b/tests/test_pooling_params.py @@ -52,13 +52,19 @@ class MockModelConfig: def test_removed_pooling_parameters(parameter: str, value: Any, message: str): data = {"input": "hello", parameter: value} for request_type in (EmbeddingRequest, ClassificationRequest, PoolingRequest): - with pytest.raises(ValidationError, match=message) as exc_info: + with pytest.raises(VLLMValidationError, match=message): TypeAdapter(request_type).validate_python(data) - assert len(exc_info.value.errors()) == 1 - with pytest.raises(ValidationError, match=message) as exc_info: - TypeAdapter(PoolerConfig).validate_python({parameter: value}) - assert len(exc_info.value.errors()) == 1 + # PoolerConfig still raises bare ValueError for `normalize` + # (wrapped to ValidationError by Pydantic), but `check_removed_pooling_task` + # raises VLLMValidationError for removed tasks. + if parameter == "normalize": + with pytest.raises(ValidationError, match=message) as exc_info: + TypeAdapter(PoolerConfig).validate_python({parameter: value}) + assert len(exc_info.value.errors()) == 1 + else: + with pytest.raises(VLLMValidationError, match=message): + TypeAdapter(PoolerConfig).validate_python({parameter: value}) if parameter == "task": with pytest.raises(VLLMValidationError, match=message): @@ -80,7 +86,7 @@ def test_embed(): invalid_parameters = classify_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(embed_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -100,7 +106,7 @@ def test_embed_dimensions(model_info: EmbedModelInfo): pooling_params = PoolingParams(task=task, dimensions=None) pooling_params.verify(model_config) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, dimensions=1) pooling_params.verify(model_config) @@ -131,7 +137,7 @@ def test_embed_dimensions_matryoshka_without_list_upper_bound(): PoolingParams(task=task, dimensions=16).verify(model_config) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): PoolingParams(task=task, dimensions=64).verify(model_config) @@ -150,7 +156,7 @@ def test_classify(task): invalid_parameters = embed_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(classify_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -176,7 +182,7 @@ def test_token_embed(pooling_type: str): invalid_parameters = classify_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(embed_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) @@ -202,6 +208,6 @@ def test_token_classify(pooling_type: str): invalid_parameters = embed_parameters + step_pooling_parameters for p in set(invalid_parameters) - set(classify_parameters): - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): pooling_params = PoolingParams(task=task, **{p: True}) pooling_params.verify(model_config) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py index e5d811fbb13..65ab0738c96 100644 --- a/tests/test_sampling_params.py +++ b/tests/test_sampling_params.py @@ -5,6 +5,7 @@ from dataclasses import dataclass import pytest from vllm import SamplingParams +from vllm.exceptions import VLLMValidationError @dataclass @@ -32,7 +33,7 @@ class MockModelConfig: ) def test_diffusion_rejects_unsupported_params(kwargs: dict): params = SamplingParams(**kwargs) - with pytest.raises(ValueError, match="not yet supported with diffusion"): + with pytest.raises(VLLMValidationError, match="not yet supported with diffusion"): params.verify(MockModelConfig(is_diffusion=True), None, None, None) diff --git a/tests/tool_use/test_chat_completion_request_validations.py b/tests/tool_use/test_chat_completion_request_validations.py index 7adf4beb9d8..1def2acc56c 100644 --- a/tests/tool_use/test_chat_completion_request_validations.py +++ b/tests/tool_use/test_chat_completion_request_validations.py @@ -4,6 +4,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.exceptions import VLLMValidationError def test_chat_completion_request_with_no_tools(): @@ -27,7 +28,7 @@ def test_chat_completion_request_with_no_tools(): assert request.tool_choice == "none" # tools key present but empty -- should be rejected - with pytest.raises(ValueError, match="must not be an empty array"): + with pytest.raises(VLLMValidationError, match="must not be an empty array"): ChatCompletionRequest.model_validate( { "messages": [{"role": "user", "content": "Hello"}], @@ -40,7 +41,7 @@ def test_chat_completion_request_with_no_tools(): @pytest.mark.parametrize("tool_choice", ["auto", "required"]) def test_chat_completion_request_with_tool_choice_but_no_tools(tool_choice): with pytest.raises( - ValueError, match="When using `tool_choice`, `tools` must be set." + VLLMValidationError, match="When using `tool_choice`, `tools` must be set." ): ChatCompletionRequest.model_validate( { @@ -51,7 +52,7 @@ def test_chat_completion_request_with_tool_choice_but_no_tools(tool_choice): ) with pytest.raises( - ValueError, match="When using `tool_choice`, `tools` must be set." + VLLMValidationError, match="When using `tool_choice`, `tools` must be set." ): ChatCompletionRequest.model_validate( { @@ -134,7 +135,7 @@ SAMPLE_TOOL = { def test_structured_outputs_with_named_tool_choice_rejected(): """structured_outputs cannot be combined with a named tool_choice.""" with pytest.raises( - ValueError, + VLLMValidationError, match="structured outputs or tools, not both", ): ChatCompletionRequest.model_validate( @@ -168,7 +169,7 @@ def test_structured_outputs_with_auto_tool_choice_allowed(): def test_multiple_structured_outputs_rejected(): """Only one kind of structured output constraint is allowed.""" with pytest.raises( - ValueError, + VLLMValidationError, match="You can only use one kind of constraints", ): ChatCompletionRequest.model_validate( diff --git a/tests/tool_use/test_responses_request_validations.py b/tests/tool_use/test_responses_request_validations.py index 59b156b76a1..b50482960c8 100644 --- a/tests/tool_use/test_responses_request_validations.py +++ b/tests/tool_use/test_responses_request_validations.py @@ -2,12 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from pydantic import ValidationError from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ResponsesResponse, ) +from vllm.exceptions import VLLMValidationError SAMPLE_TOOL = { "type": "function", @@ -58,13 +58,13 @@ def test_responses_request_required_without_tools(tools): if tools is not None: kwargs["tools"] = tools with pytest.raises( - ValidationError, match="Tool choice 'required' must be specified" + VLLMValidationError, match="Tool choice 'required' must be specified" ): ResponsesRequest.model_validate(kwargs) def test_responses_request_named_tool_choice_without_tools(): - with pytest.raises(ValidationError, match="not found in 'tools' parameter"): + with pytest.raises(VLLMValidationError, match="not found in 'tools' parameter"): ResponsesRequest.model_validate( { "input": "Hello", @@ -107,7 +107,7 @@ def test_responses_request_named_tool_choice_matching(): def test_responses_request_named_tool_choice_not_matching(): - with pytest.raises(ValidationError, match="not found in 'tools' parameter"): + with pytest.raises(VLLMValidationError, match="not found in 'tools' parameter"): ResponsesRequest.model_validate( { "input": "Hello", @@ -164,7 +164,7 @@ def test_responses_request_empty_tools_tool_choice_auto(): ], ) def test_responses_request_named_tool_choice_missing_name(tool_choice): - with pytest.raises(ValidationError, match="not found in 'tools' parameter"): + with pytest.raises(VLLMValidationError, match="not found in 'tools' parameter"): ResponsesRequest.model_validate( { "input": "Hello", @@ -176,7 +176,7 @@ def test_responses_request_named_tool_choice_missing_name(tool_choice): def test_responses_request_empty_tools_named_tool_choice(): - with pytest.raises(ValidationError, match="not found in 'tools' parameter"): + with pytest.raises(VLLMValidationError, match="not found in 'tools' parameter"): ResponsesRequest.model_validate( { "input": "Hello", diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index cd0aff79de8..955e5c5bf02 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -59,7 +59,7 @@ def test_decoder_max_context_length_validation( "Make sure that `max_model_len` is no smaller than the number of " "text tokens (prompt + requested output tokens)." ) - with pytest.raises(ValueError) as excinfo: + with pytest.raises(VLLMValidationError) as excinfo: vllm_model.generate_greedy(prompt_ids, max_tokens) assert expected_msg in str(excinfo.value) diff --git a/tests/v1/e2e/general/test_min_tokens.py b/tests/v1/e2e/general/test_min_tokens.py index bb041cd3862..c5b6341fb98 100644 --- a/tests/v1/e2e/general/test_min_tokens.py +++ b/tests/v1/e2e/general/test_min_tokens.py @@ -16,6 +16,7 @@ Covers: import pytest from vllm import LLM, SamplingParams +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput # Test configuration @@ -479,13 +480,13 @@ def test_min_tokens_validation(): # Invalid cases with pytest.raises( - ValueError, + VLLMValidationError, match="min_tokens must be greater than or equal to 0", ): SamplingParams(min_tokens=-1, max_tokens=10) with pytest.raises( - ValueError, + VLLMValidationError, match="min_tokens must be less than or equal to max_tokens", ): SamplingParams(min_tokens=15, max_tokens=10) diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py index 01c5fe6f8eb..1954ce6a7dc 100644 --- a/tests/v1/e2e/general/test_streaming_input.py +++ b/tests/v1/e2e/general/test_streaming_input.py @@ -20,6 +20,7 @@ import pytest_asyncio from vllm import SamplingParams from vllm.engine.protocol import StreamingInput +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind @@ -571,13 +572,17 @@ async def test_streaming_input_validation_errors(engine: AsyncLLM): yield StreamingInput(prompt="test") # Test n > 1 is rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_n2 = SamplingParams(max_tokens=10, n=2) async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"): pass # Test FINAL_ONLY is rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_final = SamplingParams( max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY ) @@ -585,7 +590,9 @@ async def test_streaming_input_validation_errors(engine: AsyncLLM): pass # Test stop strings are rejected - with pytest.raises(ValueError, match="Input streaming not currently supported"): + with pytest.raises( + VLLMValidationError, match="Input streaming not currently supported" + ): params_stop = SamplingParams(max_tokens=10, stop=["stop"]) async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"): pass diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index afb6e4c98b7..3845d6297a7 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -19,6 +19,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.exceptions import VLLMValidationError from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform @@ -485,7 +486,7 @@ async def test_dp_rank_argument(): pass # Test with out-of-range DP rank. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): async for _ in engine.generate( request_id="request-35", prompt=TEXT_PROMPT, @@ -554,8 +555,8 @@ async def test_header_dp_rank_argument(): # Test 2: Out-of-range DP rank (1) mock_raw_request.headers = {"X-data-parallel-rank": "1"} - # should raise ValueError for out-of-range rank - with pytest.raises(ValueError): + # should raise VLLMValidationError for out-of-range rank + with pytest.raises(VLLMValidationError): await serving_chat.create_chat_completion(req, mock_raw_request) diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index fc8ce50c05f..f57ea285eb7 100644 --- a/tests/v1/logits_processors/utils.py +++ b/tests/v1/logits_processors/utils.py @@ -11,6 +11,7 @@ import torch from tests.utils import requires_spawn_multiprocessing from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.v1.sample.logits_processor import ( @@ -61,7 +62,7 @@ class DummyLogitsProcessor(LogitsProcessor): "target_token" ) if target_token is not None and not isinstance(target_token, int): - raise ValueError( + raise VLLMValidationError( f"target_token value {target_token} {type(target_token)} is not int" ) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index aa17d2a1004..d643b0b6fde 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -405,7 +405,7 @@ def test_max_logprobs(): runner.generate(["Hello world"], sampling_params=vllm_sampling_params) bad_sampling_params = SamplingParams(logprobs=2) - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): runner.generate(["Hello world"], sampling_params=bad_sampling_params) diff --git a/tests/v1/sample/test_sampling_params_e2e.py b/tests/v1/sample/test_sampling_params_e2e.py index 56b93ea1e01..d385e96b7a2 100644 --- a/tests/v1/sample/test_sampling_params_e2e.py +++ b/tests/v1/sample/test_sampling_params_e2e.py @@ -4,6 +4,7 @@ import pytest from vllm import LLM, SamplingParams +from vllm.exceptions import VLLMValidationError MODEL = "hmellor/tiny-random-LlamaForCausalLM" PROMPT = "Hello my name is Robert and I" @@ -161,15 +162,15 @@ def test_allowed_token_ids(llm): assert output[0].outputs[0].token_ids[-1] == token_id # Reject empty allowed_token_ids. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[])) # Reject negative token id. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[-1])) # Reject out of vocabulary. - with pytest.raises(ValueError): + with pytest.raises(VLLMValidationError): _ = llm.generate(PROMPT, SamplingParams(allowed_token_ids=[10000000])) diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 31ce961ff61..7ea60cc6609 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,6 +5,7 @@ import pytest from vllm.config import StructuredOutputsConfig +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams, StructuredOutputsParams pytestmark = pytest.mark.cpu_test @@ -32,7 +33,7 @@ def test_structured_outputs_rejected_for_diffusion_models(): params = SamplingParams( structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) ) - with pytest.raises(ValueError, match="not yet supported for diffusion"): + with pytest.raises(VLLMValidationError, match="not yet supported for diffusion"): params._validate_structured_outputs( _StubModelConfig(is_diffusion=True), StructuredOutputsConfig(), @@ -63,7 +64,7 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match): rejected at request validation (-> 400) instead of reaching and crashing the engine.""" params = SamplingParams(structured_outputs=structured_outputs) - with pytest.raises(ValueError, match=match): + with pytest.raises(VLLMValidationError, match=match): params._validate_structured_outputs( _StubModelConfig(is_diffusion=False), StructuredOutputsConfig(), diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 9103dd7fae9..f57e320a906 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -28,7 +28,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware @@ -42,20 +41,15 @@ from vllm.entrypoints.serve.utils.api_utils import ( ) from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.server_utils import ( - engine_error_handler, exception_handler, - generation_error_handler, get_uvicorn_log_config, http_exception_handler, lifespan, log_response, validation_exception_handler, + vllm_error_handler, ) -from vllm.exceptions import ( - VLLMNotFoundError, - VLLMUnprocessableEntityError, - VLLMValidationError, -) +from vllm.exceptions import VLLMError from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_derenderer import OnlineDerenderer @@ -67,7 +61,6 @@ from vllm.usage.usage_lib import UsageContext from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import is_valid_ipv6_address from vllm.utils.system_utils import decorate_logs, set_ulimit -from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.version import __version__ as VLLM_VERSION prometheus_multiproc_dir: tempfile.TemporaryDirectory @@ -291,23 +284,26 @@ def build_app( allow_headers=args.allowed_headers, ) + # Exception handlers are registered in four layers: + # 1. framework errors raised by FastAPI/Starlette + # 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler + # 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` + # 4. the raw ``Exception`` handler as a safety net + # Registering specific exception types (rather than only ``Exception``) + # ensures they are handled by ``ExceptionMiddleware`` (inside the Prometheus + # middleware) rather than ``ServerErrorMiddleware`` (outside it), so their + # status codes are recorded correctly. app.exception_handler(HTTPException)(http_exception_handler) app.exception_handler(RequestValidationError)(validation_exception_handler) - app.exception_handler(EngineGenerateError)(engine_error_handler) - app.exception_handler(EngineDeadError)(engine_error_handler) - app.exception_handler(GenerationError)(generation_error_handler) - # Register specific exception types so they are handled by - # ExceptionMiddleware (inside the Prometheus middleware) rather than - # ServerErrorMiddleware (outside it). Without this, these exceptions - # propagate through Prometheus as unhandled and get recorded as 5xx - # even though they result in 4xx responses to the client. - app.exception_handler(VLLMValidationError)(exception_handler) - app.exception_handler(VLLMUnprocessableEntityError)(exception_handler) - app.exception_handler(VLLMNotFoundError)(exception_handler) + + app.exception_handler(VLLMError)(vllm_error_handler) + + # TODO(zqzten): remove these fallback handlers after migration to VLLMError app.exception_handler(ValueError)(exception_handler) app.exception_handler(TypeError)(exception_handler) app.exception_handler(OverflowError)(exception_handler) app.exception_handler(NotImplementedError)(exception_handler) + app.exception_handler(Exception)(exception_handler) # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 05536f14217..810804d7631 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -19,7 +19,7 @@ from pydantic import ( from vllm.config.utils import replace from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMServerError, VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import StructuredOutputsParams from vllm.utils import random_uuid @@ -404,7 +404,7 @@ class DeltaMessage(OpenAIBaseModel): return data -class GenerationError(Exception): +class GenerationError(VLLMServerError): """raised when finish_reason indicates internal server error (500)""" def __init__(self, message: str = "Internal server error"): diff --git a/vllm/entrypoints/serve/utils/error_response.py b/vllm/entrypoints/serve/utils/error_response.py index fc17a75c75a..2aa785c53bb 100644 --- a/vllm/entrypoints/serve/utils/error_response.py +++ b/vllm/entrypoints/serve/utils/error_response.py @@ -28,7 +28,9 @@ def create_error_response( ) from vllm.exceptions import ( + VLLMClientError, VLLMNotFoundError, + VLLMServerError, VLLMUnprocessableEntityError, VLLMValidationError, ) @@ -45,8 +47,23 @@ def create_error_response( err_type = "NotFoundError" status_code = HTTPStatus.NOT_FOUND param = None + elif isinstance(exc, VLLMClientError): + # Any other client-caused error defaults to 400. + err_type = "BadRequestError" + status_code = HTTPStatus.BAD_REQUEST + param = None + elif isinstance(exc, GenerationError): + err_type = "InternalServerError" + status_code = exc.status_code + param = None + elif isinstance(exc, VLLMServerError): + # Any other server-caused error defaults to 500. + err_type = "InternalServerError" + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + param = None + # Fallback for raw exceptions not yet migrated to VLLMError. + # TODO(zqzten): remove these fallback handlers after migration to VLLMError elif isinstance(exc, (ValueError, TypeError, OverflowError)): - # Common validation errors from user input err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = None @@ -54,10 +71,6 @@ def create_error_response( err_type = "NotImplementedError" status_code = HTTPStatus.NOT_IMPLEMENTED param = None - elif isinstance(exc, GenerationError): - err_type = "InternalServerError" - status_code = exc.status_code - param = None elif any(cls.__name__ == "TemplateError" for cls in type(exc).__mro__): # jinja2.TemplateError and its subclasses (avoid importing jinja2) err_type = "BadRequestError" diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index c6520658090..c60abfae9a2 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -31,7 +31,7 @@ from vllm.entrypoints.serve.utils.error_response import ( create_error_response, sanitize_message, ) -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMError, VLLMValidationError from vllm.logger import init_logger from vllm.utils.gc_utils import freeze_gc_heap from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError @@ -325,6 +325,16 @@ async def log_response(request: Request, call_next): return response +async def vllm_error_handler(req: Request, exc: VLLMError): + """Dispatch a vLLM-specific error to the appropriate handler.""" + if isinstance(exc, (EngineGenerateError, EngineDeadError)): + return await engine_error_handler(req, exc) + elif isinstance(exc, GenerationError): + return await generation_error_handler(req, exc) + else: + return await exception_handler(req, exc) + + async def engine_error_handler( req: Request, exc: EngineDeadError | EngineGenerateError ): diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 4112c3de24b..4383e9a6441 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -6,7 +6,25 @@ from typing import Any -class VLLMValidationError(ValueError): +class VLLMError(Exception): + """Base class for all vLLM-specific errors. + + Subclasses are split into `VLLMClientError` (caused by the request, mapped + to 4xx) and `VLLMServerError` (caused by the server, mapped to 5xx). + Dispatching on this hierarchy lets the entrypoints decide the HTTP status + without relying on raw Python exception types such as `ValueError`. + """ + + +class VLLMClientError(VLLMError): + """Base class for errors caused by the client request (4xx).""" + + +class VLLMServerError(VLLMError): + """Base class for errors caused by the server (5xx).""" + + +class VLLMValidationError(VLLMClientError): """vLLM-specific validation error for request validation failures. Args: @@ -36,7 +54,7 @@ class VLLMValidationError(ValueError): return f"{base} ({', '.join(extras)})" if extras else base -class VLLMNotFoundError(Exception): +class VLLMNotFoundError(VLLMClientError): """vLLM-specific NotFoundError""" pass @@ -66,7 +84,7 @@ class LoRAAdapterNotFoundError(VLLMNotFoundError): return self.message -class VLLMUnprocessableEntityError(ValueError): +class VLLMUnprocessableEntityError(VLLMClientError): """vLLM-specific error for unprocessable entity requests. This exception is raised when the request content is invalid or cannot be diff --git a/vllm/inputs/engine.py b/vllm/inputs/engine.py index f997004d2fb..bda40bebbc8 100644 --- a/vllm/inputs/engine.py +++ b/vllm/inputs/engine.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Literal, TypeAlias from typing_extensions import NotRequired, TypedDict, assert_never +from vllm.exceptions import VLLMValidationError + if TYPE_CHECKING: import torch @@ -284,7 +286,7 @@ which can be passed to `LLMEngine.add_request` or `AsyncLLM.add_request`. def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: if enc_input["type"] == "embeds": - raise ValueError( + raise VLLMValidationError( "Embedding inputs are not supported for encoder-decoder models" ) @@ -302,7 +304,7 @@ def _validate_enc_input(enc_input: SingletonInput) -> EncoderInput: def _validate_dec_input(dec_input: SingletonInput) -> DecoderEngineInput: if dec_input["type"] == "embeds": - raise ValueError( + raise VLLMValidationError( "Embedding inputs are not supported for encoder-decoder models" ) diff --git a/vllm/pooling_params.py b/vllm/pooling_params.py index 6cb130fdbbb..5280e997c9c 100644 --- a/vllm/pooling_params.py +++ b/vllm/pooling_params.py @@ -7,6 +7,7 @@ from typing import Any import msgspec from vllm.config import ModelConfig, PoolerConfig +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind from vllm.tasks import PoolingTask, check_removed_pooling_task @@ -145,7 +146,7 @@ class PoolingParams( invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise VLLMValidationError( f"Task {self.task} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -170,21 +171,21 @@ class PoolingParams( valid_range = f"[1, {embedding_size}]" dimensions_in_range = 1 <= dimensions <= embedding_size if not model_config.is_matryoshka: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} does not support Matryoshka " f"embeddings; dimensions must be unset " f"(received dimensions={dimensions})." ) if not dimensions_in_range: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} only supports dimensions in " f"range {valid_range}, got {dimensions}." ) mds = model_config.matryoshka_dimensions if mds is not None and dimensions not in mds: - raise ValueError( + raise VLLMValidationError( f"Model {model_name!r} only supports Matryoshka " f"dimensions {str(mds)}, got {dimensions}." ) @@ -208,7 +209,7 @@ class PoolingParams( invalid_parameters.append(k) if invalid_parameters: - raise ValueError( + raise VLLMValidationError( f"Task {self.task!r} only supports {valid_parameters} " f"parameters, does not support " f"{invalid_parameters} parameters" @@ -231,7 +232,7 @@ class PoolingParams( def __post_init__(self) -> None: check_removed_pooling_task(self.task) if self.output_kind != RequestOutputKind.FINAL_ONLY: - raise ValueError( + raise VLLMValidationError( "For pooling output_kind has to be FINAL_ONLY, " f"got {self.output_kind!r}" ) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 25e36ceb568..5dedbde372e 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -100,12 +100,12 @@ class StructuredOutputsParams: ] ) if count > 1: - raise ValueError( + raise VLLMValidationError( "You can only use one kind of structured outputs constraint " f"but multiple are specified: {self.__dict__}" ) if count < 1: - raise ValueError( + raise VLLMValidationError( "You must use one kind of structured outputs constraint " f"but none are specified: {self.__dict__}" ) @@ -166,13 +166,13 @@ class RepetitionDetectionParams: or self.min_pattern_size < 0 or self.min_pattern_size > self.max_pattern_size ): - raise ValueError( + raise VLLMValidationError( "max_pattern_size, min_pattern_size must be >=0, " "with min_pattern_size <= max_pattern_size. " "Set both to 0 to disable repetitive pattern detection." ) if self.max_pattern_size > 0 and self.min_count < 2: - raise ValueError( + raise VLLMValidationError( "min_count must be >= 2 to detect repetitive patterns " "in engine output. If you do not wish to detect repetitive " "patterns, set max_pattern_size to 0." @@ -514,31 +514,33 @@ class SamplingParams( def _verify_args(self) -> None: if not isinstance(self.n, int): - raise ValueError(f"n must be an int, but is of type {type(self.n)}") + raise VLLMValidationError( + f"n must be an int, but is of type {type(self.n)}" + ) if self.n < 1: - raise ValueError(f"n must be at least 1, got {self.n}.") + raise VLLMValidationError(f"n must be at least 1, got {self.n}.") max_n = envs.VLLM_MAX_N_SEQUENCES if self.n > max_n: - raise ValueError( + raise VLLMValidationError( f"n must be at most {max_n}, got {self.n}. " "To increase this limit, set the VLLM_MAX_N_SEQUENCES " "environment variable." ) if not -2.0 <= self.presence_penalty <= 2.0: - raise ValueError( + raise VLLMValidationError( f"presence_penalty must be in [-2, 2], got {self.presence_penalty}." ) if not -2.0 <= self.frequency_penalty <= 2.0: - raise ValueError( + raise VLLMValidationError( f"frequency_penalty must be in [-2, 2], got {self.frequency_penalty}." ) if not math.isfinite(self.repetition_penalty): - raise ValueError( + raise VLLMValidationError( "repetition_penalty must be a finite number, " f"got {self.repetition_penalty}." ) if self.repetition_penalty <= 0.0: - raise ValueError( + raise VLLMValidationError( "repetition_penalty must be greater than zero, got " f"{self.repetition_penalty}." ) @@ -568,15 +570,15 @@ class SamplingParams( ) # quietly accept -1 as disabled, but prefer 0 if self.top_k < -1: - raise ValueError( + raise VLLMValidationError( f"top_k must be 0 (disable), or at least 1, got {self.top_k}." ) if not isinstance(self.top_k, int): - raise TypeError( + raise VLLMValidationError( f"top_k must be an integer, got {type(self.top_k).__name__}" ) if not 0.0 <= self.min_p <= 1.0: - raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.") + raise VLLMValidationError(f"min_p must be in [0, 1], got {self.min_p}.") if self.max_tokens is not None and self.max_tokens < 1: raise VLLMValidationError( f"max_tokens must be at least 1, got {self.max_tokens}.", @@ -584,11 +586,11 @@ class SamplingParams( value=self.max_tokens, ) if self.min_tokens < 0: - raise ValueError( + raise VLLMValidationError( f"min_tokens must be greater than or equal to 0, got {self.min_tokens}." ) if self.max_tokens is not None and self.min_tokens > self.max_tokens: - raise ValueError( + raise VLLMValidationError( f"min_tokens must be less than or equal to " f"max_tokens={self.max_tokens}, got {self.min_tokens}." ) @@ -617,27 +619,29 @@ class SamplingParams( ) assert isinstance(self.stop_token_ids, list) if not all(isinstance(st_id, int) for st_id in self.stop_token_ids): - raise ValueError( + raise VLLMValidationError( f"stop_token_ids must contain only integers, got {self.stop_token_ids}." ) assert isinstance(self.stop, list) if any(not stop_str for stop_str in self.stop): - raise ValueError("stop cannot contain an empty string.") + raise VLLMValidationError("stop cannot contain an empty string.") if self.stop and not self.detokenize: - raise ValueError( + raise VLLMValidationError( "stop strings are only supported when detokenize is True. " "Set detokenize=True to use stop." ) assert isinstance(self.bad_words, list) if any(not bad_word for bad_word in self.bad_words): - raise ValueError( + raise VLLMValidationError( f"bad_words cannot contain an empty string. " f"Got bad_words={self.bad_words}" ) def _verify_greedy_sampling(self) -> None: if self.n > 1: - raise ValueError(f"n must be 1 when using greedy sampling, got {self.n}.") + raise VLLMValidationError( + f"n must be 1 when using greedy sampling, got {self.n}." + ) def update_from_generation_config( self, @@ -889,7 +893,7 @@ class SamplingParams( # Some sampling parameters are not yet compatible with spec decoding. if self.min_p > _SAMPLING_EPS or self.logit_bias: - raise ValueError( + raise VLLMValidationError( "The min_p and logit_bias sampling parameters " "are not yet supported with speculative decoding." ) @@ -910,7 +914,7 @@ class SamplingParams( or self.bad_words or self.allowed_token_ids ): - raise ValueError( + raise VLLMValidationError( "The temperature, min_p, seed, min_tokens, logit_bias, " "bad_words, and allowed_token_ids sampling parameters " "are not yet supported with diffusion models." @@ -930,7 +934,7 @@ class SamplingParams( # rather than sampling left-to-right, which the grammar FSM # requires. Without this check, requests fail mid-generation # with an FSM rejection (HTTP 500). See issue #45436. - raise ValueError( + raise VLLMValidationError( "Structured outputs are not yet supported for diffusion " "language models. Remove the structured output constraint " "(e.g. `response_format`, `structured_outputs`) from the " @@ -938,7 +942,7 @@ class SamplingParams( ) if tokenizer is None: - raise ValueError( + raise VLLMValidationError( "Structured outputs requires a tokenizer so it can't be used with 'skip_tokenizer_init'" # noqa: E501 ) @@ -952,7 +956,7 @@ class SamplingParams( if backend != _backend and not ( backend == "auto" and self.structured_outputs._backend_was_auto ): - raise ValueError( + raise VLLMValidationError( "Request-level structured output backend selection is not " f"supported. The request specified '{_backend}', but vLLM " f"was initialised with '{backend}'. This error can be " @@ -967,7 +971,7 @@ class SamplingParams( and not self.structured_outputs.choice ): # It is invalid for choice to be an empty list - raise ValueError( + raise VLLMValidationError( f"Choice '{self.structured_outputs.choice}' cannot be an empty list" # noqa: E501 ) # Reject empty string grammar early to avoid engine-side crashes @@ -975,16 +979,20 @@ class SamplingParams( isinstance(self.structured_outputs.grammar, str) and self.structured_outputs.grammar.strip() == "" ): - raise ValueError("structured_outputs.grammar cannot be an empty string") + raise VLLMValidationError( + "structured_outputs.grammar cannot be an empty string" + ) # Reject empty string json schema early to avoid engine-side crashes if ( isinstance(self.structured_outputs.json, str) and self.structured_outputs.json.strip() == "" ): - raise ValueError("structured_outputs.json cannot be an empty string") + raise VLLMValidationError( + "structured_outputs.json cannot be an empty string" + ) # Reject json_object=False early to avoid engine-side crashes if self.structured_outputs.json_object is False: - raise ValueError( + raise VLLMValidationError( "structured_outputs.json_object must be True if set; omit " "structured_outputs to disable structured outputs" ) @@ -1006,7 +1014,7 @@ class SamplingParams( validate_xgrammar_grammar(self) elif backend.startswith("guidance"): if _is_non_tekken_mistral(tokenizer=tokenizer): - raise ValueError( + raise VLLMValidationError( "Non-tekken Mistral tokenizers are not supported for the 'guidance'" " structured output backend. Please either use a more recent " "Mistral model, the ['xgrammar', 'outlines'] " @@ -1026,7 +1034,7 @@ class SamplingParams( elif backend == "lm-format-enforcer": # lm format enforcer backend if is_mistral_tokenizer(tokenizer): - raise ValueError( + raise VLLMValidationError( "Mistral tokenizer is not supported for the 'lm-format-enforcer' " "structured output backend. Please use ['xgrammar', 'outlines'] " "backends or tokenizer_mode='hf' instead." diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index e8d33c961dc..1718a5adc39 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -21,6 +21,7 @@ from vllm.distributed.weight_transfer.base import ( from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient, StreamingInput from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep +from vllm.exceptions import VLLMClientError, VLLMValidationError from vllm.inputs import EngineInput, PromptType from vllm.logger import init_logger from vllm.lora.request import LoRARequest @@ -309,7 +310,7 @@ class AsyncLLM(EngineClient): and not is_pooling and params.prompt_logprobs ): - raise ValueError( + raise VLLMValidationError( "--kv-sharing-fast-prefill produces incorrect logprobs for " "prompt tokens, please disable it when the requests need " "prompt logprobs" @@ -476,7 +477,7 @@ class AsyncLLM(EngineClient): ) req.external_req_id = request_id if req.prompt_embeds is not None: - raise ValueError( + raise VLLMValidationError( "prompt_embeds not supported for streaming inputs" ) prompt_text, _, _ = extract_prompt_components( @@ -512,7 +513,7 @@ class AsyncLLM(EngineClient): or params.output_kind == RequestOutputKind.FINAL_ONLY or params.stop ): - raise ValueError( + raise VLLMValidationError( "Input streaming not currently supported " "for pooling models, n > 1, request_kind = FINAL_ONLY " "or with stop strings." @@ -604,7 +605,7 @@ class AsyncLLM(EngineClient): raise # Request validation error. - except ValueError as e: + except VLLMClientError as e: if self.log_requests: logger.info("Request %s failed (bad request): %s.", request_id, e) raise @@ -869,7 +870,7 @@ class AsyncLLM(EngineClient): raise # Request validation error. - except ValueError: + except VLLMClientError: if self.log_requests: logger.info("Request %s failed (bad request).", request_id) raise diff --git a/vllm/v1/engine/exceptions.py b/vllm/v1/engine/exceptions.py index d9f79a019e2..edb0fe4261b 100644 --- a/vllm/v1/engine/exceptions.py +++ b/vllm/v1/engine/exceptions.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -class EngineGenerateError(Exception): +from vllm.exceptions import VLLMServerError + + +class EngineGenerateError(VLLMServerError): """Raised when a AsyncLLM.generate() fails. Recoverable.""" pass -class EngineDeadError(Exception): +class EngineDeadError(VLLMServerError): """Raised when the EngineCore dies. Unrecoverable.""" def __init__(self, *args, suppress_context: bool = False, **kwargs): diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 10a6329f5a0..735593672da 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -7,6 +7,7 @@ from typing import Any, Literal import vllm.envs as envs from vllm.config import VllmConfig +from vllm.exceptions import VLLMValidationError from vllm.inputs import ( EngineInput, PromptType, @@ -90,7 +91,7 @@ class InputProcessor: task for task in supported_tasks if task in GENERATION_TASKS ] if not supported_generation_tasks: - raise ValueError("This model does not support generation") + raise VLLMValidationError("This model does not support generation") params.verify( self.model_config, @@ -104,13 +105,13 @@ class InputProcessor: self.vllm_config.reasoning_config is None or not self.vllm_config.reasoning_config.enabled ): - raise ValueError( + raise VLLMValidationError( "thinking_token_budget is set but reasoning_config is " "not configured. Please set --reasoning-parser " "and/or --reasoning-config to use thinking_token_budget." ) if self.use_v2_model_runner: - raise ValueError( + raise VLLMValidationError( "thinking_token_budget is not yet supported by the V2 " "model runner. Run vLLM with VLLM_USE_V2_MODEL_RUNNER=0 " "to use thinking_token_budget." @@ -120,7 +121,7 @@ class InputProcessor: task for task in supported_tasks if task in POOLING_TASKS ] if not supported_pooling_tasks: - raise ValueError("This model does not support pooling") + raise VLLMValidationError("This model does not support pooling") if params.task is None: if "token_embed" in supported_pooling_tasks: @@ -131,7 +132,7 @@ class InputProcessor: params.task = "plugin" if params.task not in supported_pooling_tasks: - raise ValueError( + raise VLLMValidationError( f"Unsupported task: {params.task!r} " f"Supported tasks: {supported_pooling_tasks}" ) @@ -149,7 +150,7 @@ class InputProcessor: # LoRA request passed in while LoRA is not enabled if not self.lora_config: - raise ValueError( + raise VLLMValidationError( f"Got lora_request {lora_request} but LoRA is not enabled!" ) @@ -261,7 +262,7 @@ class InputProcessor: dp_local_size = parallel_config.data_parallel_size_local num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks): - raise ValueError( + raise VLLMValidationError( f"data_parallel_rank {data_parallel_rank} " f"is out of range [0, {num_ranks})." ) @@ -393,7 +394,7 @@ class InputProcessor: return if prompt_len == 0 and prompt_type == "decoder": - raise ValueError(f"The {prompt_type} prompt cannot be empty") + raise VLLMValidationError(f"The {prompt_type} prompt cannot be empty") model_config = self.model_config max_prompt_len = ( @@ -415,7 +416,7 @@ class InputProcessor: "number of text tokens." ) - raise ValueError( + raise VLLMValidationError( f"The {prompt_type} prompt (length {prompt_len}) is " f"longer than the maximum model length of {max_prompt_len}. " f"{suggestion}" @@ -425,7 +426,7 @@ class InputProcessor: "Make sure that `max_model_len` is no smaller than the " "number of text tokens (prompt + requested output tokens)." ) - raise ValueError( + raise VLLMValidationError( f"The {prompt_type} prompt (length {prompt_len}) plus the number of " f"requested output tokens (at least 1) is longer than the maximum " f"model length of {max_prompt_len}. {suggestion}" @@ -457,7 +458,7 @@ class InputProcessor: for mm_position in mm_positions: num_embeds = mm_position.get_num_embeds() if num_embeds > self.mm_encoder_cache_size: - raise ValueError( + raise VLLMValidationError( f"The {prompt_type} prompt contains a(n) {modality} item " f"with {num_embeds} embedding tokens, which exceeds the " f"pre-allocated encoder cache size " @@ -481,7 +482,9 @@ class InputProcessor: # truly out-of-vocabulary. model_vocab_size = model_config.get_vocab_size() if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1): - raise ValueError(f"Token id {max_input_id} is out of vocabulary") + raise VLLMValidationError( + f"Token id {max_input_id} is out of vocabulary" + ) def _validate_model_inputs( self, diff --git a/vllm/v1/sample/logits_processor/__init__.py b/vllm/v1/sample/logits_processor/__init__.py index 2cb89e1ea95..f319097495d 100644 --- a/vllm/v1/sample/logits_processor/__init__.py +++ b/vllm/v1/sample/logits_processor/__init__.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING import torch +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.logits_process import LogitsProcessor as RequestLogitsProcessor from vllm.sampling_params import SamplingParams @@ -228,7 +229,12 @@ def validate_logits_processors_parameters( tuple(logits_processors) if logits_processors is not None else None ) for logits_procs in cached_load_custom_logitsprocs(logits_processors): - logits_procs.validate_params(sampling_params) + try: + logits_procs.validate_params(sampling_params) + except ValueError as e: + # Legacy custom logitsprocs may still raise ValueError from + # validate_params; convert for backward compatibility. + raise VLLMValidationError(str(e)) from e class AdapterLogitsProcessor(LogitsProcessor): diff --git a/vllm/v1/sample/logits_processor/interface.py b/vllm/v1/sample/logits_processor/interface.py index 3e426e321b3..dce0706335e 100644 --- a/vllm/v1/sample/logits_processor/interface.py +++ b/vllm/v1/sample/logits_processor/interface.py @@ -62,7 +62,9 @@ class LogitsProcessor(ABC): def validate_params(cls, sampling_params: SamplingParams): """Validate sampling params for this logits processor. - Raise ValueError for invalid ones. + Raise ``VLLMValidationError`` (preferred) / ``ValueError`` (backward compatible) + for invalid params. Bare ``ValueError`` is converted to ``VLLMValidationError`` + at the engine boundary so online serving returns HTTP 400. """ return None From 0bb548b60e88b74304008817b4908b65ec8b2393 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 28 Jul 2026 23:27:45 -0500 Subject: [PATCH 61/67] [CI][ROCm] Stabilize Qwen2-VL LoRA test (#50161) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .buildkite/test_areas/lora.yaml | 1 - tests/lora/test_qwenvl.py | 22 +++++++++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 214d4a12bf7..a79196ffbd8 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -16,7 +16,6 @@ steps: amd: dind: false device: mi300_1 - soft_fail: true working_dir: "/vllm-workspace/tests" timeout_in_minutes: 85 source_file_dependencies: diff --git a/tests/lora/test_qwenvl.py b/tests/lora/test_qwenvl.py index 3cbb534bdad..3362aa46ed2 100644 --- a/tests/lora/test_qwenvl.py +++ b/tests/lora/test_qwenvl.py @@ -186,6 +186,18 @@ QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct" QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct" +def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: + # These tests assert exact greedy outputs. Force the Triton LoRA shrink + # kernel to use SPLIT_K=1 so it stores the complete reduction directly + # instead of accumulating split-K partial results with atomic_add. This + # targets reduction determinism, not full batch invariance. + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + # The kernel configuration reads VLLM_BATCH_INVARIANT at import time. + # Spawn the engine process so it observes this setting even if the LoRA + # Triton utilities were already imported during test collection. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + def test_qwen2vl_lora(qwen2vl_lora_files): """Test Qwen 2.0 VL model with LoRA""" config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files) @@ -250,7 +262,12 @@ def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files): ) -def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files): +def test_qwen3vl_vision_lora( + qwen3vl_vision_lora_files, + monkeypatch: pytest.MonkeyPatch, +): + _enable_deterministic_lora_shrink(monkeypatch) + config = TestConfig( model_path=QWEN3VL_MODEL_PATH, lora_path=qwen3vl_vision_lora_files, @@ -273,6 +290,7 @@ def test_qwen2vl_multiple_lora_types( qwen2vl_language_lora_files, qwen2vl_vision_tower_connector_lora_files, qwen2vl_vision_tower_lora_files, + monkeypatch: pytest.MonkeyPatch, ): """ Test multiple LoRA adapter types (language, vision tower + connector, @@ -283,6 +301,8 @@ def test_qwen2vl_multiple_lora_types( the multimodal encoder cache correctly manages state transitions between language-only and vision-enabled LoRA adapters. """ + _enable_deterministic_lora_shrink(monkeypatch) + config = TestConfig( model_path=QWEN2VL_MODEL_PATH, # We'll override the lora_path for each specific test, but need to provide From dc1be79031d948d7a18c37600881e45ca708d913 Mon Sep 17 00:00:00 2001 From: Philip Pesic <116986036+philippesic@users.noreply.github.com> Date: Wed, 29 Jul 2026 04:34:55 +0000 Subject: [PATCH 62/67] Add CachePolicyFactory for pluggable/external eviction policies (#49114) Signed-off-by: Philip Pesic --- docs/features/kv_offloading_usage.md | 33 +++++- tests/v1/kv_offload/cpu/__init__.py | 0 tests/v1/kv_offload/cpu/policies/__init__.py | 0 .../kv_offload/cpu/policies/test_factory.py | 111 ++++++++++++++++++ tests/v1/kv_offload/cpu/test_manager.py | 2 + vllm/v1/kv_offload/base.py | 8 -- vllm/v1/kv_offload/cpu/manager.py | 25 ++-- vllm/v1/kv_offload/cpu/policies/arc.py | 2 +- vllm/v1/kv_offload/cpu/policies/base.py | 4 +- vllm/v1/kv_offload/cpu/policies/factory.py | 87 ++++++++++++++ vllm/v1/kv_offload/cpu/policies/lru.py | 1 + vllm/v1/kv_offload/cpu/spec.py | 6 +- vllm/v1/kv_offload/factory.py | 7 ++ vllm/v1/kv_offload/tiering/manager.py | 4 +- vllm/v1/kv_offload/tiering/spec.py | 12 +- 15 files changed, 269 insertions(+), 33 deletions(-) create mode 100644 tests/v1/kv_offload/cpu/__init__.py create mode 100644 tests/v1/kv_offload/cpu/policies/__init__.py create mode 100644 tests/v1/kv_offload/cpu/policies/test_factory.py create mode 100644 vllm/v1/kv_offload/cpu/policies/factory.py diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 13dbd5299d3..72f838c7d2c 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -70,7 +70,8 @@ vllm serve \ | `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). | | `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. | | `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. | -| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. | +| `eviction_policy` | no | `lru` | both | Primary tier policy: built-in `lru`/`arc`, or a custom `CachePolicy` name (see [Custom Eviction Policies](#custom-eviction-policies)). | +| `cache_policy_module_path` | no | — | both | Python import path for a custom `CachePolicy` not in the built-in registry. Required only when `eviction_policy` is not built-in and wasn't pre-registered via `CachePolicyFactory` (advanced). | | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. | | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | @@ -78,6 +79,36 @@ vllm serve \ | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | +## Custom Eviction Policies + +`eviction_policy` resolves through `CachePolicyFactory` (`vllm/v1/kv_offload/cpu/policies/factory.py`), which pre-registers the built-in `lru` and `arc` policies. + +### Out-of-tree (recommended) + +Implement `CachePolicy` (`vllm/v1/kv_offload/cpu/policies/base.py`) in your own package — no vLLM fork or patch required — and point `kv_connector_extra_config` at it directly: + +```json +{ + "cpu_bytes_to_use": 10737418240, + "eviction_policy": "MyCachePolicy", + "cache_policy_module_path": "my_package.my_module" +} +``` + +`eviction_policy` is checked against the built-in registry first; if it isn't a registered name, vLLM imports `cache_policy_module_path` and looks up `eviction_policy` as a class name in that module — the same fallback `spec_module_path` provides for a custom `OffloadingSpec`. No import or registration call needs to run before the server starts. + +### Registering a friendly short name (in-process only) + +If you control the process that constructs the vLLM engine (e.g. an embedding application), you can register a short name once at startup instead of repeating the module path in every config: + +```python +from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory + +CachePolicyFactory.register_cache_policy("my_policy", "my_package.my_module", "MyCachePolicy") +``` + +Then set `"eviction_policy": "my_policy"` in `kv_connector_extra_config`, the same as `"lru"`/`"arc"`. This only takes effect within the process that ran the `register_cache_policy` call — it does not help when the server is launched as a separate process (e.g. via the `vllm serve` CLI), where the out-of-tree `cache_policy_module_path` config above is the only option. + ## Secondary Tiers Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields. diff --git a/tests/v1/kv_offload/cpu/__init__.py b/tests/v1/kv_offload/cpu/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/v1/kv_offload/cpu/policies/__init__.py b/tests/v1/kv_offload/cpu/policies/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/v1/kv_offload/cpu/policies/test_factory.py b/tests/v1/kv_offload/cpu/policies/test_factory.py new file mode 100644 index 00000000000..14ccf8b67e7 --- /dev/null +++ b/tests/v1/kv_offload/cpu/policies/test_factory.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable + +import pytest + +from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager +from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory +from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy + + +class _DummyCachePolicy(CachePolicy): + """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded + by module path, so it must be importable at module scope (mirrors + tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec).""" + + def __init__(self, cache_capacity: int) -> None: + self.cache_capacity = cache_capacity + + def get(self, key: OffloadKey) -> BlockStatus | None: + return None + + def insert(self, key: OffloadKey, block: BlockStatus) -> None: + pass + + def remove(self, key: OffloadKey) -> None: + pass + + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + pass + + def evict( + self, n: int, protected: set[OffloadKey] + ) -> list[tuple[OffloadKey, BlockStatus]] | None: + return None + + def clear(self) -> None: + pass + + +@pytest.fixture(autouse=True) +def restore_cache_policy_registry(): + """Save and restore CachePolicyFactory._registry between tests.""" + original = dict(CachePolicyFactory._registry) + yield + CachePolicyFactory._registry = original + + +class TestCachePolicyFactory: + """Unit tests for CachePolicyFactory (registration/resolution by name).""" + + def test_pre_registered_policies_can_be_imported(self): + """If someone moves a policy module but forgets to update + factory.py, CI fails.""" + for name in CachePolicyFactory._registry: + cls = CachePolicyFactory._registry[name]() + assert issubclass(cls, CachePolicy) + + def test_lru_and_arc_registered(self): + assert CachePolicyFactory.get_cache_policy_cls("lru") is LRUCachePolicy + assert CachePolicyFactory.get_cache_policy_cls("arc") is ARCCachePolicy + + def test_register_and_resolve_custom_policy(self): + CachePolicyFactory.register_cache_policy( + "dummy", + "tests.v1.kv_offload.cpu.policies.test_factory", + "_DummyCachePolicy", + ) + policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy") + assert policy_cls is _DummyCachePolicy + + manager = CPUOffloadingManager(num_blocks=4, cache_policy="dummy") + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_raises(self): + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent") + + def test_duplicate_registration_raises(self): + with pytest.raises(ValueError, match="is already registered"): + CachePolicyFactory.register_cache_policy("lru", "some.module", "SomeClass") + + def test_dynamic_load_via_cache_policy_module_path(self): + """Out-of-tree policy loaded via cache_policy_module_path, no + register_cache_policy() call -- this is how external projects + integrate a custom CachePolicy without forking/patching vLLM. + Mirrors tests/v1/kv_offload/test_factory.py's + test_dynamic_load_via_spec_module_path.""" + policy_cls = CachePolicyFactory.get_cache_policy_cls( + "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory" + ) + assert policy_cls is _DummyCachePolicy + + def test_manager_resolves_policy_via_module_path(self): + """End-to-end: CPUOffloadingManager resolves an unregistered policy + purely from cache_policy_module_path.""" + manager = CPUOffloadingManager( + num_blocks=4, + cache_policy="_DummyCachePolicy", + cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory", + ) + assert isinstance(manager._policy, _DummyCachePolicy) + + def test_unregistered_policy_without_module_path_raises(self): + """eviction_policy not in registry + no cache_policy_module_path -> + ValueError, same shape as the OffloadingSpecFactory error path.""" + with pytest.raises(ValueError, match="Unknown cache policy"): + CachePolicyFactory.get_cache_policy_cls("nonexistent", None) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 6520a93fd1a..5a36be16bf2 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -37,6 +37,7 @@ _EMPTY_REQ_CTX = make_req_context() def make_cpu_manager( num_blocks: int = 4, cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, store_threshold: int = 0, max_tracker_size: int = 64_000, @@ -44,6 +45,7 @@ def make_cpu_manager( return CPUOffloadingManager( num_blocks=num_blocks, cache_policy=cache_policy, + cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 6ff09f23d37..b0db86cf9bf 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -13,8 +13,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar import numpy as np import torch -from vllm.logger import init_logger - if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, @@ -27,8 +25,6 @@ from vllm.v1.kv_offload.config import OffloadingConfig # Use the helper functions below to construct / decompose keys. OffloadKey = NewType("OffloadKey", bytes) -logger = init_logger(__name__) - def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey: """Pack a block hash and group index into an `OffloadKey`.""" @@ -540,10 +536,6 @@ class OffloadingSpec(ABC): return {} def __init__(self, config: OffloadingConfig): - logger.warning( - "Initializing OffloadingSpec. This API is experimental and " - "subject to change in the future as we iterate the design." - ) self.config = config self.extra_config = config.extra_config self.replicated_layout: bool = False diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index c2ec4170b8e..0eef5cf7e79 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict from collections.abc import Collection, Iterable -from typing import Literal from typing_extensions import override @@ -24,19 +23,15 @@ from vllm.v1.kv_offload.cpu.common import ( CPULoadStoreSpec, CPUOffloadingMetrics, ) -from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy -from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy - -_CACHE_POLICIES: dict[str, type[CachePolicy]] = { - "lru": LRUCachePolicy, - "arc": ARCCachePolicy, -} +from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory class CPUOffloadingManager(OffloadingManager): """ - An OffloadingManager with a pluggable CachePolicy (LRU or ARC). + An OffloadingManager with a pluggable CachePolicy, resolved by name via + CachePolicyFactory (built in: "lru", "arc"; external policies can either + register their own or be loaded out-of-tree via cache_policy_module_path). The manager owns all shared logic: ref-counting, event emission, block pool management, and the prepare_store/complete_store skeletons. @@ -47,7 +42,8 @@ class CPUOffloadingManager(OffloadingManager): def __init__( self, num_blocks: int, - cache_policy: Literal["lru", "arc"] = "lru", + cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, store_threshold: int = 1, max_tracker_size: int = 64_000, @@ -57,12 +53,9 @@ class CPUOffloadingManager(OffloadingManager): self._num_allocated_blocks: int = 0 self._free_list: list[int] = [] self.events: list[OffloadingEvent] | None = [] if enable_events else None - policy_cls = _CACHE_POLICIES.get(cache_policy) - if policy_cls is None: - raise ValueError( - f"Unknown cache policy: {cache_policy!r}. " - f"Supported: {list(_CACHE_POLICIES)}" - ) + policy_cls = CachePolicyFactory.get_cache_policy_cls( + cache_policy, cache_policy_module_path + ) self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks) # Track the number of blocks in the cache that are evictable. i.e. ref_cnt 0. self._num_evictable_cache_blocks: int = 0 diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index f682a47e45f..d6569cbcfd2 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -48,7 +48,7 @@ class ARCCachePolicy(CachePolicy): """ def __init__(self, cache_capacity: int): - self.cache_capacity: int = cache_capacity + super().__init__(cache_capacity) self.target_t1_size: float = 0.0 self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index 2b6681e4992..908907d326f 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -41,8 +41,8 @@ class CachePolicy(ABC): and eviction, so they cannot be separated cleanly. """ - @abstractmethod - def __init__(self, cache_capacity: int) -> None: ... + def __init__(self, cache_capacity: int) -> None: + self.cache_capacity = cache_capacity @abstractmethod def get(self, key: OffloadKey) -> BlockStatus | None: diff --git a/vllm/v1/kv_offload/cpu/policies/factory.py b/vllm/v1/kv_offload/cpu/policies/factory.py new file mode 100644 index 00000000000..88d11b270c9 --- /dev/null +++ b/vllm/v1/kv_offload/cpu/policies/factory.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +from collections.abc import Callable + +from vllm.logger import init_logger +from vllm.v1.kv_offload.cpu.policies.base import CachePolicy + +logger = init_logger(__name__) + + +class CachePolicyFactory: + """Registry for CachePolicy implementations, resolved by name. + + Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in + policies are pre-registered below. External policies can either + register_cache_policy() a friendly short name up front, or skip + registration entirely and pass a module path at lookup time (out-of-tree, + no vLLM fork/patch required) -- see get_cache_policy_cls. + """ + + _registry: dict[str, Callable[[], type[CachePolicy]]] = {} + + @classmethod + def register_cache_policy( + cls, name: str, module_path: str, class_name: str + ) -> None: + """Register a cache policy with a lazy-loading module and class name.""" + if name in cls._registry: + raise ValueError(f"Cache policy '{name}' is already registered.") + + def loader() -> type[CachePolicy]: + module = importlib.import_module(module_path) + return getattr(module, class_name) + + cls._registry[name] = loader + + @classmethod + def get_cache_policy_cls( + cls, name: str, module_path: str | None = None + ) -> type[CachePolicy]: + """Get a cache policy class by name. + + Args: + name: Name of the cache policy. Checked against the registry + first; if it's not registered and `module_path` is given, + `name` is imported from there instead -- an out-of-tree + policy needs no register_cache_policy() call at all, just + this module path passed through config (mirrors + OffloadingSpecFactory.get_spec_cls's spec_module_path + fallback). + module_path: Python import path to load `name` from when it is + not a registered policy. + + Returns: + The cache policy class. + + Raises ValueError if the cache policy is neither registered nor + resolvable via `module_path`. + """ + if name in cls._registry: + return cls._registry[name]() + if module_path is None: + raise ValueError( + f"Unknown cache policy: {name!r}. Supported: {list(cls._registry)}. " + "For an out-of-tree policy, also set cache_policy_module_path." + ) + logger.warning( + "Loading out-of-tree cache policy '%s' from '%s'. This API is " + "experimental and subject to change in the future as we " + "iterate the design.", + name, + module_path, + ) + module = importlib.import_module(module_path) + policy_cls = getattr(module, name) + assert issubclass(policy_cls, CachePolicy) + return policy_cls + + +# Register built-in policies here. +CachePolicyFactory.register_cache_policy( + "lru", "vllm.v1.kv_offload.cpu.policies.lru", "LRUCachePolicy" +) +CachePolicyFactory.register_cache_policy( + "arc", "vllm.v1.kv_offload.cpu.policies.arc", "ARCCachePolicy" +) diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index efa24fe9033..e8ccf0bdef5 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -19,6 +19,7 @@ class LRUCachePolicy(CachePolicy): """ def __init__(self, cache_capacity: int): + super().__init__(cache_capacity) # Blocks with ref_cnt 0 (not participating in any loads/stores) ordered in LRU self.evictable_blocks: OrderedDict[OffloadKey, None] = OrderedDict() self.blocks: dict[OffloadKey, BlockStatus] = {} diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index d755bfecbc4..9162c7b18b4 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -117,6 +117,9 @@ class CPUOffloadingSpec(OffloadingSpec): self._worker: CPUOffloadingWorker | None = None self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru") + self.cache_policy_module_path: str | None = self.extra_config.get( + "cache_policy_module_path" + ) @override def get_manager(self) -> OffloadingManager: @@ -131,7 +134,8 @@ class CPUOffloadingSpec(OffloadingSpec): self._manager = CPUOffloadingManager( num_blocks=self.num_blocks, - cache_policy=self.eviction_policy, # type: ignore[arg-type] + cache_policy=self.eviction_policy, + cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, store_threshold=store_threshold, max_tracker_size=max_tracker_size, diff --git a/vllm/v1/kv_offload/factory.py b/vllm/v1/kv_offload/factory.py index 931fda8308f..19bc401277d 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -35,6 +35,13 @@ class OffloadingSpecFactory: spec_module_path = extra_config.get("spec_module_path") if spec_module_path is None: raise ValueError(f"Unsupported spec type: {spec_name}") + logger.warning( + "Loading out-of-tree offloading spec '%s' from '%s'. This " + "API is experimental and subject to change in the future " + "as we iterate the design.", + spec_name, + spec_module_path, + ) spec_module = importlib.import_module(spec_module_path) spec_cls = getattr(spec_module, spec_name) assert issubclass(spec_cls, OffloadingSpec) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index c6738096135..d1d3c421159 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -93,11 +93,13 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): num_blocks: int, mmap_region: SharedOffloadRegion, cache_policy: str = "lru", + cache_policy_module_path: str | None = None, enable_events: bool = False, ): super().__init__( num_blocks=num_blocks, - cache_policy=cache_policy, # type: ignore[arg-type] + cache_policy=cache_policy, + cache_policy_module_path=cache_policy_module_path, enable_events=enable_events, ) self._mmap_region = mmap_region diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 0d200a42390..e6ec192415e 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -9,8 +9,13 @@ and configurable secondary tiers (e.g., Storage, Network). Configuration via kv_connector_extra_config: - cpu_bytes_to_use: (required) Bytes to allocate for CPU primary tier - block_size: (optional) Block size for offloaded blocks (default: GPU block size) - - eviction_policy: (optional) Primary tier eviction policy: "lru" or - "arc" (default: "lru") + - eviction_policy: (optional) Primary tier eviction policy: built-in "lru"/ + "arc", or the name of a policy registered via CachePolicyFactory, or an + out-of-tree CachePolicy class name paired with cache_policy_module_path + (default: "lru") + - cache_policy_module_path: (optional) Python import path to load + eviction_policy from when it names an out-of-tree CachePolicy not + registered via CachePolicyFactory - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - type: (required) Type of secondary tier (e.g., "example", "storage", "network") @@ -178,7 +183,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec): # Create primary tier (CPU-based) primary_tier = CPUPrimaryTierOffloadingManager( num_blocks=self.num_blocks, - cache_policy=self.eviction_policy, # type: ignore[arg-type] + cache_policy=self.eviction_policy, + cache_policy_module_path=self.cache_policy_module_path, enable_events=self.kv_events_config.enable_kv_cache_events, mmap_region=scheduler_mmap, ) From db7a79cbf74ec0b21e009faa7f1b563f1f88aeb2 Mon Sep 17 00:00:00 2001 From: Rehan Khan Date: Wed, 29 Jul 2026 10:22:17 +0530 Subject: [PATCH 63/67] [CPU] Fix s390x builds and update torch version in dockerfile (#50144) Signed-off-by: Rehan Khan Co-authored-by: Li, Jiang --- csrc/cpu/cpu_types_vxe.hpp | 6 +++--- docker/Dockerfile.s390x | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index bf96554a8df..26a14e8ea6b 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -269,7 +269,7 @@ struct FP32Vec4 : public Vec { explicit FP32Vec4(__vector float data) : reg(data) {} - explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} + FP32Vec4(const FP32Vec4& data) : reg(data.reg) {} }; struct FP32Vec8 : public Vec { @@ -298,7 +298,7 @@ struct FP32Vec8 : public Vec { explicit FP32Vec8(f32x4x2_t data) : reg(data) {} - explicit FP32Vec8(const FP32Vec8& data) { + FP32Vec8(const FP32Vec8& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; } @@ -643,7 +643,7 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(f32x4x4_t data) : reg(data) {} - explicit FP32Vec16(const FP32Vec16& data) { + FP32Vec16(const FP32Vec16& data) { reg.val[0] = data.reg.val[0]; reg.val[1] = data.reg.val[1]; reg.val[2] = data.reg.val[2]; diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index b71e035e152..d645e37e78c 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -61,13 +61,13 @@ ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" FROM python-install AS torch-vision # Install torchvision -ARG TORCH_VISION_VERSION=v0.26.0 +ARG TORCH_VISION_VERSION=v0.28.0 WORKDIR /tmp RUN --mount=type=cache,target=/root/.cache/uv \ git clone https://github.com/pytorch/vision.git && \ cd vision && \ git checkout $TORCH_VISION_VERSION && \ - uv pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu && \ + uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu && \ python setup.py bdist_wheel FROM python-install AS hf-xet-builder From 6f91edf96d3f3272945809c04702380053bff4de Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Wed, 29 Jul 2026 12:54:57 +0800 Subject: [PATCH 64/67] [Test] dynamic_shapes_compilation (#49974) Signed-off-by: JaredforReal Co-authored-by: Isotr0py --- .../test_dynamic_shapes_compilation.py | 100 ++++++++---------- 1 file changed, 45 insertions(+), 55 deletions(-) diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 7f725c14f21..3260d5aecc5 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import gc import tempfile from contextlib import contextmanager @@ -9,8 +8,7 @@ import pytest import torch from tests.models.utils import check_logprobs_close -from tests.utils import wait_for_rocm_memory_to_settle -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.compilation.decorators import support_torch_compile from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config from vllm.config.compilation import ( @@ -49,6 +47,7 @@ def get_test_models(): @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_dynamic_shapes_compilation( monkeypatch, + vllm_runner, model_name, shapes_type, use_aot_compile, @@ -79,9 +78,13 @@ def test_dynamic_shapes_compilation( print(f"Testing {shapes_type.name} dynamic shapes...") - # Initialize the model with specific dynamic shapes configuration - model = LLM( - model=model_name, + sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) + test_prompts = [prompt, "The capital of France is"] + + # VllmRunner shuts down the engine core on exit, so the eager model + # below never races a lingering compiled engine for GPU memory. + with vllm_runner( + model_name, compilation_config={ "mode": CompilationMode.VLLM_COMPILE, "dynamic_shapes_config": { @@ -90,33 +93,25 @@ def test_dynamic_shapes_compilation( }, }, max_model_len=1024, - ) + enable_chunked_prefill=None, + ) as vllm_model: + compiled_outputs = [] + for p in test_prompts: + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Compiled model produced empty output" + compiled_outputs.append((output.token_ids, output.text, output.logprobs)) - sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10) - test_prompts = [prompt, "The capital of France is"] - - compiled_outputs = [] - for p in test_prompts: - output = model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Compiled model produced empty output" - compiled_outputs.append((output.token_ids, output.text, output.logprobs)) - - del model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - wait_for_rocm_memory_to_settle() - - eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024) - eager_outputs = [] - for p in test_prompts: - output = eager_model.generate(p, sampling_params)[0].outputs[0] - assert len(output.text.strip()) > 0, "Eager model produced empty output" - eager_outputs.append((output.token_ids, output.text, output.logprobs)) - del eager_model - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + with vllm_runner( + model_name, + enforce_eager=True, + max_model_len=1024, + enable_chunked_prefill=None, + ) as vllm_model: + eager_outputs = [] + for p in test_prompts: + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0] + assert len(output.text.strip()) > 0, "Eager model produced empty output" + eager_outputs.append((output.token_ids, output.text, output.logprobs)) check_logprobs_close( outputs_0_lst=eager_outputs, @@ -241,44 +236,39 @@ def test_model_specialization_with_evaluate_guards( @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") -def test_piecewise_backend_empty_sym_shape_indices(): +def test_piecewise_backend_empty_sym_shape_indices(vllm_runner): """Test that PiecewiseBackend handles empty sym_shape_indices correctly. When all inputs have static shapes (no torch.SymInt), sym_shape_indices will be empty. The fix in PiecewiseBackend.__call__ handles this case by using the first compiled range_entry. """ - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() - # Use small max_model_len and max_num_batched_tokens to encourage # static shape compilation with empty sym_shape_indices - llm = LLM( - model="Qwen/Qwen3-0.6B", + with vllm_runner( + "Qwen/Qwen3-0.6B", max_model_len=512, max_num_batched_tokens=1, + enable_chunked_prefill=None, compilation_config={ "mode": CompilationMode.VLLM_COMPILE, "dynamic_shapes_config": { "type": DynamicShapesType.BACKED.value, }, }, - ) + ) as vllm_model: + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) + # Generate with static shape inputs + output = vllm_model.llm.generate( + "Hello, my name is", sampling_params=sampling_params + ) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output" - # Generate with static shape inputs - output = llm.generate("Hello, my name is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output" - - # Generate again to verify compilation works with empty sym_shape_indices - output = llm.generate("The capital of France is", sampling_params=sampling_params) - result = output[0].outputs[0].text - assert len(result) > 0, "Should generate non-empty output on second run" - - del llm - gc.collect() - torch.accelerator.empty_cache() - torch.accelerator.synchronize() + # Generate again to verify compilation works with empty sym_shape_indices + output = vllm_model.llm.generate( + "The capital of France is", sampling_params=sampling_params + ) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output on second run" From 6f00a1ae3bd4b86168667bce673998218f461c0f Mon Sep 17 00:00:00 2001 From: Artur Fierka Date: Wed, 29 Jul 2026 07:15:10 +0200 Subject: [PATCH 65/67] fused_moe: add VLLM_TRITON_USE_TD tensor-descriptor path (#42436) Signed-off-by: Artur Fierka Signed-off-by: Lena Onyshchenko <162571002+oonyshch@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Lena Onyshchenko <162571002+oonyshch@users.noreply.github.com> --- tests/kernels/moe/test_moe.py | 31 +++++-- .../layers/fused_moe/fused_moe.py | 81 +++++++++++++++---- .../layers/fused_moe/oracle/unquantized.py | 7 ++ vllm/model_executor/layers/fused_moe/utils.py | 81 +++++++++++++++++++ 4 files changed, 179 insertions(+), 21 deletions(-) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 9c43aa97409..a519980dead 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -36,6 +36,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( batched_fused_marlin_moe, fused_marlin_moe, ) +from vllm.model_executor.layers.fused_moe.utils import ( + moe_use_td_hw_supported, +) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_permute_bias, ) @@ -53,9 +56,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types +from vllm.triton_utils import tl from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + def iterative_moe( hidden_states: torch.Tensor, @@ -289,6 +295,7 @@ def run_moe_test( @pytest.mark.parametrize("ep_size", EP_SIZE) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("padding", [True, False]) +@pytest.mark.parametrize("use_td", [False, True]) def test_fused_moe( m: int, n: int, @@ -298,9 +305,19 @@ def test_fused_moe( ep_size: int, dtype: torch.dtype, padding: bool, + use_td: bool, monkeypatch, workspace_init, ): + if use_td and not hasattr(tl, "make_tensor_descriptor"): + pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor") + if use_td and not moe_use_td_hw_supported(): + pytest.skip( + "tensor_descriptor.gather requires XPU or NVIDIA Blackwell " + "(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas " + "rejects on Hopper (sm90) and earlier" + ) + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") set_random_seed(7) # @@ -311,17 +328,17 @@ def test_fused_moe( # Setup test data # - a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 - w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 - w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 + a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 + w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 - score = torch.randn((m, e), device="cuda", dtype=dtype) + score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype) if ep_size > 1: local_e = e // ep_size - e_ids = torch.randint(0, e, (local_e,), device="cuda", dtype=torch.int32) - e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32) - e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32) + e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32) + e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32) + e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32) w1 = w1[e_ids] w2 = w2[e_ids] else: diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 0d357fcbf45..be4930052a9 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -28,9 +28,12 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( from vllm.model_executor.layers.fused_moe.utils import ( enable_swap_ab, moe_kernel_quantize_input, + resolve_moe_use_td, + warn_if_moe_use_td_ineffective, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.triton_utils.allocation import set_triton_allocator from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op @@ -347,6 +350,8 @@ def fused_moe_kernel( per_channel_quant: tl.constexpr, HAS_BIAS: tl.constexpr, SWAP_AB: tl.constexpr, + # Tensor-descriptor path for the A gather and B load in the K-loop. + USE_TD: tl.constexpr = False, ): """ Implements the fused computation for a Mixture of Experts (MOE) using @@ -436,7 +441,25 @@ def fused_moe_kernel( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) - if SWAP_AB: + # TD gather and the SWAP_AB accumulator layout are mutually exclusive. + tl.static_assert(not (USE_TD and SWAP_AB)) + if USE_TD: + # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx. + m_td = num_valid_tokens // top_k + a_desc = tl.make_tensor_descriptor( + base=a_ptr, + shape=(m_td, K), + strides=(stride_am, stride_ak), + block_shape=(1, BLOCK_SIZE_K), + ) + b_desc = tl.make_tensor_descriptor( + base=b_ptr + off_experts * stride_be, + shape=(N, K), + strides=(stride_bn, stride_bk), + block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K), + ) + gather_idx = (offs_token // top_k).to(tl.int32) + elif SWAP_AB: a_ptrs = a_ptr + ( offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am ) @@ -454,7 +477,6 @@ def fused_moe_kernel( + off_experts * stride_be + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) ) - if use_int8_w8a16: b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn @@ -498,18 +520,21 @@ def fused_moe_kernel( for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. - if SWAP_AB: + if USE_TD: + a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K) + b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T + elif SWAP_AB: a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + a = tl.load(a_ptrs, mask=a_mask, other=0.0) + b = tl.load(b_ptrs, mask=b_mask, other=0.0) else: - a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K) - b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K - a = tl.load( - a_ptrs, - mask=a_mask, - other=0.0, - ) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -536,9 +561,10 @@ def fused_moe_kernel( accumulator += tl.dot(a, b) else: accumulator += tl.dot(a, b) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk + if not USE_TD: + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk if SWAP_AB: accumulator = tl.trans(accumulator, (1, 0)) @@ -765,6 +791,19 @@ def invoke_fused_moe_triton_kernel( else: SWAP_AB = False + # Quantized weights always carry a B_scale (see the asserts below); key off + # that rather than enumerating quant flags, which misses w8a16-fp8/nvfp4/etc. + is_quantized = B_scale is not None + warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized) + + # TD path is unvalidated under quantization; fall back to the pointer path. + use_td = resolve_moe_use_td() and not is_quantized + if use_td: + # The TD path builds a tensor descriptor inside the kernel, which + # requires a PyTorch-backed scratch allocator to be registered + # (Triton raises "no allocator was set" otherwise on CUDA). + set_triton_allocator(A.device) + if use_fp8_w8a8 or use_int8_w8a8: assert B_scale is not None assert block_shape is None or triton.cdiv( @@ -805,6 +844,19 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + if use_td and A.size(1) % BLOCK_SIZE_K != 0: + # TD gather/load feeding tl.dot with a non-block-aligned K + # miscompiles (~74% of output elements wrong) on real HW; + # this is a compiler-codegen issue, not a Python-maskable + # boundary gap. Fall back to the pointer-arith path. + logger.warning_once( + "Disabling VLLM_TRITON_USE_TD for this MoE launch: K=%d is not " + "a multiple of BLOCK_SIZE_K=%d, which triggers a known " + "Triton tensor-descriptor + tl.dot miscompilation.", + A.size(1), + BLOCK_SIZE_K, + ) + use_td = False fused_moe_kernel[grid]( A, B, @@ -847,6 +899,7 @@ def invoke_fused_moe_triton_kernel( HAS_BIAS=HAS_BIAS, BLOCK_SIZE_K=BLOCK_SIZE_K, SWAP_AB=SWAP_AB, + USE_TD=use_td, **config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 8fa1a0c265c..bf77a316bb4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -362,6 +362,13 @@ def make_unquantized_moe_kernel( experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: + from vllm.model_executor.layers.fused_moe.utils import ( + warn_if_moe_use_td_ineffective, + ) + + # Warn against the selected backend, not each probed candidate. + warn_if_moe_use_td_ineffective(backend.value, is_quantized=False) + # Create Prepare/Finalize is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) prepare_finalize = maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index e3d6493dda2..cce8ccd073f 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -7,7 +7,9 @@ from typing import TYPE_CHECKING import torch import torch.nn.functional as F +import vllm.envs as envs from vllm import _custom_ops as ops +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -39,6 +41,8 @@ from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +logger = init_logger(__name__) + @triton.jit def _count_expert_num_tokens( @@ -585,3 +589,80 @@ def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool: and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N >= 64 ) + + +def moe_use_td_hw_supported() -> bool: + """Whether the current device can run the TD (gather) path of + ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override). + + The A-load uses ``tensor_descriptor.gather``, which lowers to the PTX + ``tile::gather4`` instruction. That instruction is part of the + ``tcgen05``/Tensor Memory (TMEM) family introduced with Blackwell and has + no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature + '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike + ``scatter4``, ``gather4`` is supported across the whole sm100+ range + including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498, + which enables ``gather4`` on sm120/sm121 while leaving ``scatter4`` + unsupported there. So this gates on a blanket ``has_device_capability(100)`` + rather than the sm100 *family* check used for the scatter store path. + """ + if current_platform.is_xpu(): + return True + if current_platform.is_cuda(): + return current_platform.has_device_capability(100) + return False + + +def resolve_moe_use_td() -> bool: + """Tri-state resolver for ``VLLM_TRITON_USE_TD``. + + Unset auto-selects the TD path on XPU only, mirroring the attention + dispatcher in ``triton_attn.py``. ``1``/``0`` force it on/off regardless + of hardware; forcing ``1`` where it cannot compile (see + ``moe_use_td_hw_supported``) fails at ptxas. Blackwell CUDA (sm100+) can + compile it but is opt-in only, pending validation. + """ + override = envs.VLLM_TRITON_USE_TD + if override is None: + return current_platform.is_xpu() + return override + + +_warned_moe_use_td_ineffective = False + + +def warn_if_moe_use_td_ineffective( + active_backend: str, is_quantized: bool = False +) -> None: + """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored. + + Fires when the user set the env explicitly and either (a) the active + MoE backend is not the fused Triton kernel, or (b) the model is + quantized (the TD path falls back to the pointer path under any + quantization). + """ + global _warned_moe_use_td_ineffective + if _warned_moe_use_td_ineffective: + return + if envs.VLLM_TRITON_USE_TD is None: + return + is_triton = active_backend.upper() == "TRITON" + if is_triton and not is_quantized: + return + if not is_triton: + reason = ( + f"the active MoE backend is {active_backend!r}; pass " + "`--moe-backend triton` to enable the tensor-descriptor path" + ) + else: + reason = ( + "the model uses quantized MoE weights; the TD path is " + "currently restricted to non-quantized weights and falls " + "back to the pointer path" + ) + logger.warning( + "VLLM_TRITON_USE_TD is set to %s but %s.", + envs.VLLM_TRITON_USE_TD, + reason, + ) + _warned_moe_use_td_ineffective = True From 7c6729b769597541d327f666273cd972b8a5318d Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 29 Jul 2026 14:10:58 +0800 Subject: [PATCH 66/67] [Model] Add Kimi K3 support: model files and kernels [1/N] (#50089) Signed-off-by: zjy0516 Signed-off-by: Isotr0py Signed-off-by: Bugen Zhao Signed-off-by: Yongye Zhu Co-authored-by: Isotr0py Co-authored-by: Thien Tran Co-authored-by: Bugen Zhao Co-authored-by: Yongye Zhu Co-authored-by: Ziming Huang Co-authored-by: Roger Wang Co-authored-by: Isotr0py Co-authored-by: aoshen02 Co-authored-by: Woosuk Kwon Co-authored-by: Jee Jee Li Co-authored-by: aoshen02 Co-authored-by: Summer Yang Co-authored-by: Kevin H. Luu Co-authored-by: Bowen Wang Co-authored-by: gnovack Co-authored-by: Nick Hill Co-authored-by: Yifan Qiao Co-authored-by: xiaozhoupy Co-authored-by: Roy Wang Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Co-authored-by: Codex --- .buildkite/check-torch-abi.py | 1 + .buildkite/test_areas/kernels.yaml | 11 - .buildkite/test_areas/models_basic.yaml | 6 +- .pre-commit-config.yaml | 2 +- CMakeLists.txt | 27 +- .../kernels/benchmark_k3_cutedsl_residual.py | 367 ++++ .../benchmark_kimi_k3_latent_moe_tail.py | 806 +++++++ .../benchmark_kimi_k3_sp_collectives.py | 239 +++ cmake/external_projects/flashkda.cmake | 74 + csrc/custom_all_gather_reduce_scatter.cuh | 327 +++ csrc/custom_all_reduce.cuh | 316 +-- csrc/custom_collective_common.cuh | 332 +++ csrc/flashkda_registration.cpp | 17 + csrc/libtorch_stable/activation_kernels.cu | 108 + .../attention/merge_attn_states.cu | 39 +- csrc/libtorch_stable/cache_kernels.cu | 96 + .../custom_all_gather_reduce_scatter.cu | 362 ++++ .../custom_all_gather_reduce_scatter_ops.cpp | 29 + csrc/libtorch_stable/custom_all_reduce.cu | 8 +- csrc/libtorch_stable/dsv3_fused_a_gemm.cu | 151 +- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 1237 +++++++++++ .../kimi_k3/fused_kda_decode_kernel.cu | 1130 ++++++++++ .../moe/grouped_topk_kernels.cu | 438 +++- csrc/libtorch_stable/moe/moeTopKFuncs.cuh | 347 +-- .../moe/moe_align_sum_kernels.cu | 2 +- csrc/libtorch_stable/ops.h | 95 + csrc/libtorch_stable/torch_bindings.cpp | 92 +- .../w8a8/fp8/nvidia/quant_utils.cuh | 33 +- setup.py | 5 + tests/distributed/test_custom_all_reduce.py | 117 ++ .../test_kimi_k3_mla_fused_epilogue.py | 221 ++ tests/kernels/core/test_activation.py | 40 + tests/kernels/core/test_fused_q_kv_rmsnorm.py | 2 +- .../kernels/core/test_fused_rms_norm_gated.py | 4 +- tests/kernels/moe/test_grouped_topk.py | 266 +++ .../kernels/moe/test_moe_align_block_size.py | 1 + tests/kernels/test_bf16_skinny_gemm.py | 649 ++++++ tests/kernels/test_kda.py | 226 -- tests/models/kimi_k3/test_attn_res.py | 33 + tests/models/kimi_k3/test_eagle3.py | 130 ++ tests/models/kimi_k3/test_kda.py | 757 +++++++ tests/models/kimi_k3/test_kda_metadata.py | 411 ++++ tests/models/kimi_k3/test_latent_moe_tail.py | 145 ++ .../models/kimi_k3/test_sequence_parallel.py | 343 +++ .../multimodal/generation/test_common.py | 2 +- vllm/_custom_ops.py | 139 +- vllm/envs.py | 6 + .../kernels/linear/cute_dsl/_skinny_gemm.py | 180 ++ .../kernels/linear/cute_dsl/skinny_gemm.py | 252 +++ vllm/model_executor/layers/activation.py | 45 + .../layers/fused_moe/activation.py | 23 + .../model_executor/layers/fused_moe/config.py | 4 + .../layers/fused_moe/experts/deep_gemm_moe.py | 40 +- .../layers/fused_moe/experts/marlin_moe.py | 63 + .../fused_moe/experts/trtllm_mxfp4_moe.py | 57 +- .../fused_moe/experts/trtllm_nvfp4_moe.py | 36 +- vllm/model_executor/layers/fused_moe/layer.py | 6 + .../layers/fused_moe/modular_kernel.py | 31 +- .../fused_moe/runner/latent_moe_runner.py | 255 +++ .../layers/fused_moe/runner/moe_runner.py | 73 +- .../fused_moe/runner/moe_runner_interface.py | 1 + .../layers/mamba/ops/causal_conv1d.py | 18 +- .../layers/mamba/ops/gather_initial_states.py | 83 + vllm/models/common/__init__.py | 2 + vllm/models/common/ops/__init__.py | 9 + vllm/models/common/ops/fused_qk_rmsnorm.py | 103 + .../inkling/nvidia/ops/fa4_rel_attention.py | 11 +- vllm/models/kimi_k3/__init__.py | 26 + vllm/models/kimi_k3/amd/linear.py | 1065 ++++++++++ vllm/models/kimi_k3/amd/model.py | 249 +++ vllm/models/kimi_k3/amd/mtp.py | 403 ++++ .../kimi_k3/amd/ops/third_party/__init__.py | 2 + .../amd/ops/third_party/kda/__init__.py | 47 + .../kimi_k3/amd/ops/third_party/kda/chunk.py | 935 +++++++++ .../amd/ops/third_party/kda/chunk_intra.py | 662 ++++++ .../kda/chunk_intra_token_parallel.py | 197 ++ .../ops/third_party/kda/fused_recurrent.py | 621 ++++++ vllm/models/kimi_k3/common/__init__.py | 2 + vllm/models/kimi_k3/common/mm_preprocess.py | 330 +++ vllm/models/kimi_k3/common/mtp.py | 96 + vllm/models/kimi_k3/nvidia/dspark_mla.py | 527 +++++ vllm/models/kimi_k3/nvidia/kda.py | 775 +++++++ vllm/models/kimi_k3/nvidia/kda_metadata.py | 496 +++++ .../models/kimi_k3/nvidia/low_latency_gemm.py | 513 +++++ vllm/models/kimi_k3/nvidia/mla.py | 764 +++++++ vllm/models/kimi_k3/nvidia/model.py | 1859 +++++++++++++++++ vllm/models/kimi_k3/nvidia/mtp.py | 443 ++++ .../kimi_k3/nvidia/ops/cute_dsl/__init__.py | 2 + .../ops/cute_dsl/latent_moe_tail/__init__.py | 14 + ...educe_rmsnorm_reduce_scatter_early_exit.py | 1012 +++++++++ .../fused_add_multicast_gemm.py | 1291 ++++++++++++ .../fused_add_multicast_skinny_gemm.py | 438 ++++ .../cute_dsl/latent_moe_tail/lamport_copy.py | 226 ++ .../cute_dsl/latent_moe_tail/primitives.py | 437 ++++ .../ops/fused_mla_key_concat_kv_cache.py | 244 +++ .../kimi_k3/nvidia/ops/latent_moe_tail.py | 265 +++ .../kimi_k3/nvidia/ops/sequence_parallel.py | 73 + .../nvidia/ops/third_party/__init__.py | 2 + .../nvidia/ops/third_party/kda/__init__.py | 28 + .../nvidia/ops/third_party/kda/chunk.py | 938 +++++++++ .../nvidia/ops/third_party/kda/chunk_intra.py | 559 +++++ .../kda/chunk_intra_token_parallel.py | 178 ++ .../ops/third_party/kda/fused_recurrent.py | 671 ++++++ .../kimi_k3/nvidia/ops/vision_fa4_warmup.py | 195 ++ .../ops/fused_norm_gate.py | 412 ++++ .../flash_linear_attention/ops/kda.py | 88 +- vllm/transformers_utils/configs/kimi_k3.py | 139 ++ vllm/transformers_utils/processors/kimi_k3.py | 61 + .../attention/ops/triton_merge_attn_states.py | 54 +- 109 files changed, 28028 insertions(+), 792 deletions(-) create mode 100644 benchmarks/kernels/benchmark_k3_cutedsl_residual.py create mode 100644 benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py create mode 100644 benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py create mode 100644 cmake/external_projects/flashkda.cmake create mode 100644 csrc/custom_all_gather_reduce_scatter.cuh create mode 100644 csrc/custom_collective_common.cuh create mode 100644 csrc/flashkda_registration.cpp create mode 100644 csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu create mode 100644 csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp create mode 100644 csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu create mode 100644 csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu create mode 100644 tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py create mode 100644 tests/kernels/test_bf16_skinny_gemm.py delete mode 100644 tests/kernels/test_kda.py create mode 100644 tests/models/kimi_k3/test_eagle3.py create mode 100644 tests/models/kimi_k3/test_kda.py create mode 100644 tests/models/kimi_k3/test_kda_metadata.py create mode 100644 tests/models/kimi_k3/test_latent_moe_tail.py create mode 100644 tests/models/kimi_k3/test_sequence_parallel.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py create mode 100644 vllm/model_executor/layers/mamba/ops/gather_initial_states.py create mode 100644 vllm/models/common/__init__.py create mode 100644 vllm/models/common/ops/__init__.py create mode 100644 vllm/models/common/ops/fused_qk_rmsnorm.py create mode 100644 vllm/models/kimi_k3/amd/linear.py create mode 100644 vllm/models/kimi_k3/amd/model.py create mode 100644 vllm/models/kimi_k3/amd/mtp.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py create mode 100644 vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py create mode 100644 vllm/models/kimi_k3/common/__init__.py create mode 100644 vllm/models/kimi_k3/common/mm_preprocess.py create mode 100644 vllm/models/kimi_k3/common/mtp.py create mode 100644 vllm/models/kimi_k3/nvidia/dspark_mla.py create mode 100644 vllm/models/kimi_k3/nvidia/kda.py create mode 100644 vllm/models/kimi_k3/nvidia/kda_metadata.py create mode 100644 vllm/models/kimi_k3/nvidia/low_latency_gemm.py create mode 100644 vllm/models/kimi_k3/nvidia/mla.py create mode 100644 vllm/models/kimi_k3/nvidia/model.py create mode 100644 vllm/models/kimi_k3/nvidia/mtp.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py create mode 100644 vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py create mode 100644 vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py create mode 100644 vllm/transformers_utils/configs/kimi_k3.py create mode 100644 vllm/transformers_utils/processors/kimi_k3.py diff --git a/.buildkite/check-torch-abi.py b/.buildkite/check-torch-abi.py index 493952c33ec..cb580fb56f1 100644 --- a/.buildkite/check-torch-abi.py +++ b/.buildkite/check-torch-abi.py @@ -13,6 +13,7 @@ from torch_abi_audit.report import ExtensionReport, PackageReport # Temporary allowlist of extensions not yet on the stable ABI. # Shrink and remove over time. ALLOWED_UNSTABLE_LIBRARIES: tuple[str, ...] = ( + "_flashkda_C.abi3.so", "vllm_flash_attn/_vllm_fa2_C.abi3.so", "vllm_flash_attn/_vllm_fa3_C.abi3.so", "third_party/deep_gemm/_C*.so", diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 938f8690551..81fa9c5f1d8 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -214,17 +214,6 @@ steps: commands: - pytest -v -s kernels/mamba -- label: Kernels KDA Test - timeout_in_minutes: 25 - device: h200_18gb - source_file_dependencies: - - vllm/third_party/flash_linear_attention/ops/kda.py - - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py - - vllm/third_party/flash_linear_attention/ops/l2norm.py - - tests/kernels/test_kda.py - commands: - - pytest -v -s kernels/test_kda.py - - label: Kernels DeepGEMM Test (H100) key: kernels-deepgemm-test-h100 timeout_in_minutes: 35 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 24bfff5d756..e231e6d90f7 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -69,9 +69,11 @@ steps: - vllm/models/kimi_k3/ - csrc/libtorch_stable/kimi_k3/ - tests/models/kimi_k3/ + - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py + - tests/kernels/test_bf16_skinny_gemm.py commands: - # The native NVIDIA AttnRes kernel requires the SM100 family. - - pytest -v -s models/kimi_k3 + # The native NVIDIA Kimi K3 kernels require the SM100 family. + - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a9f744bfa5c..35b53a416de 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_install_hook_types: default_stages: - pre-commit # Run locally - manual # Run in CI -exclude: 'vllm/third_party/.*' +exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*' repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.0 diff --git a/CMakeLists.txt b/CMakeLists.txt index cdda81ea46e..3790db92631 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -421,8 +421,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu" + "csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp" "csrc/libtorch_stable/custom_all_reduce.cu" - "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu" + "csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu") if(VLLM_GPU_LANG STREQUAL "CUDA" AND DEFINED CMAKE_CUDA_COMPILER_VERSION AND @@ -1079,6 +1082,23 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") set(MLA_ARCHS) endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS + "9.0a;10.0f;12.0f" "${CUDA_ARCHS}") + endif() + if(FUSED_KDA_DECODE_ARCHS) + set(FUSED_KDA_DECODE_SRC + "csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_KDA_DECODE_SRC}" + CUDA_ARCHS "${FUSED_KDA_DECODE_ARCHS}") + set_property(SOURCE ${FUSED_KDA_DECODE_SRC} APPEND PROPERTY + COMPILE_OPTIONS "$<$:--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_KDA_DECODE_SRC}") + message(STATUS + "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS "10.0f" "${CUDA_ARCHS}") @@ -1138,6 +1158,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_COOPERATIVE_TOPK=1) endif() + if(FUSED_KDA_DECODE_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_FUSED_KDA_DECODE=1) + endif() if(KIMI_K3_ATTN_RES_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_KIMI_K3_ATTN_RES=1) @@ -1439,6 +1463,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) + include(cmake/external_projects/flashkda.cmake) include(cmake/external_projects/qutlass.cmake) include(cmake/external_projects/tml_fa4.cmake) diff --git a/benchmarks/kernels/benchmark_k3_cutedsl_residual.py b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py new file mode 100644 index 00000000000..7b9e349a792 --- /dev/null +++ b/benchmarks/kernels/benchmark_k3_cutedsl_residual.py @@ -0,0 +1,367 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi-K3 latent MoE addmm against CuTe residual GEMM. + +The benchmark covers ``BF16[M, 3584] @ BF16[7168, 3584].T + BF16[M, 7168]`` +with FP32 accumulation and BF16 output. Both backends execute through CUDA +Graph replay. Weights and residuals rotate across buffers exceeding L2 so the +comparison models the full latent MoE projection-and-add path. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import importlib.util +import json +import math +import statistics +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings import driver as cuda +from cuda.bindings.driver import CUstream +from quack.compile_utils import make_fake_tensor + +N = 7168 +K = 3584 + + +@dataclasses.dataclass(frozen=True, slots=True) +class Config: + block_size: int + outputs_per_block: int + k_unroll: int + vector_width: int = 8 + + +def parse_config(value: str) -> Config: + try: + parts = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" + ) from error + if len(parts) == 3: + return Config(*parts) + if len(parts) == 4: + return Config(*parts) + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH]" + ) + + +def production_residual_config(m: int) -> Config | None: + """The measured Latent-MoE residual config for M, from the K3 table.""" + from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + + spec = KIMI_K3_PROJECTIONS.get((N, K)) + config = spec.residual_config(m) if spec is not None else None + if config is None: + return None + return Config( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def candidate_configs(mode: str, selected: Config | None, m: int) -> list[Config]: + if mode == "selected": + if selected is not None: + return [selected] + # No explicit --config: fall back to the production table for this M. + config = production_residual_config(m) + return [config] if config is not None else [] + if mode == "baseline": + return [Config(224, 4, 2)] + return [ + Config(block_size, outputs_per_block, k_unroll, vector_width) + for vector_width in (4, 8) + for block_size in (32, 64, 128, 224, 448) + if block_size % 32 == 0 and K % (block_size * vector_width) == 0 + for outputs_per_block in (1, 2, 4, 7, 8) + if N % outputs_per_block == 0 + for k_unroll in (1, 2, 4) + ] + + +def load_kernel_class(path: Path): + spec = importlib.util.spec_from_file_location("cute_skinny_device", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load CuTe kernel from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.CuteSkinnyGemm + + +def stream() -> CUstream: + return CUstream(torch.cuda.current_stream().cuda_stream) + + +def compile_kernel(kernel_class, m: int, config: Config, max_registers: int): + element_type = cutlass.BFloat16 + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor(element_type, (m, k), divisibility=config.vector_width) + b = make_fake_tensor(element_type, (n, k), divisibility=config.vector_width) + residual = make_fake_tensor(element_type, (m, n), divisibility=1) + c = make_fake_tensor(element_type, (m, n), divisibility=1) + kernel = kernel_class( + element_type=element_type, + num_rows=m, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=True, + use_pdl=True, + ) + return cute.compile( + kernel, + a, + b, + residual, + c, + stream(), + options=( + "--enable-tvm-ffi --keep-cubin " + f"--ptxas-options -maxrregcount={max_registers} " + "--ptxas-options -lineinfo" + ), + ) + + +def resource_usage(compiled) -> dict[str, Any]: + executor = getattr(compiled, "_default_executor", None) + context = getattr(executor, "exec_context", None) + functions = getattr(context, "kernel_functions", None) + if not functions: + return {"resource_metrics_available": False} + + def attribute(name, function) -> int: + error, value = cuda.cuFuncGetAttribute(name, function) + if error != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuFuncGetAttribute failed with {error}") + return int(value) + + registers = [ + attribute(cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, function) + for function in functions + ] + local_bytes = [ + attribute( + cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + function, + ) + for function in functions + ] + return { + "resource_metrics_available": True, + "registers_per_thread": max(registers, default=0), + "spill_bytes": max(local_bytes, default=0), + } + + +def rotating_buffer_count(m: int, multiplier: float, limit: int) -> int: + properties = torch.cuda.get_device_properties(0) + bytes_per_pair = (N * K + m * N) * 2 + target = math.ceil(multiplier * properties.L2_cache_size) + return max(2, min(limit, math.ceil(target / bytes_per_pair))) + + +def graph_samples( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], None], + activation: torch.Tensor, + weights: Sequence[torch.Tensor], + residuals: Sequence[torch.Tensor], + repeats: int, + replays: int, +) -> tuple[list[float], list[torch.Tensor]]: + outputs = [torch.empty_like(residual) for residual in residuals] + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for weight, residual, output in zip(weights, residuals, outputs): + launch(activation, weight, residual, output) + for _ in range(20): + graph.replay() + torch.accelerator.synchronize() + + samples = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(replays): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / (replays * len(weights))) + return samples, outputs + + +def summarize(samples: Sequence[float]) -> dict[str, Any]: + ordered = sorted(samples) + + def percentile(fraction: float) -> float: + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + mean = statistics.mean(samples) + return { + "median_us": statistics.median(samples), + "p10_us": percentile(0.1), + "p90_us": percentile(0.9), + "mean_us": mean, + "cv_pct": statistics.pstdev(samples) / mean * 100.0, + "samples_us": list(samples), + } + + +def correctness( + output: torch.Tensor, + activation: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> dict[str, Any]: + actual = output.float() + reference = activation.float() @ weight.float().t() + residual.float() + error = (actual - reference).abs() + scaled_error = error / (reference.abs() + 1.0) + cosine = torch.nn.functional.cosine_similarity( + actual.flatten(), reference.flatten(), dim=0 + ).item() + return { + "valid": cosine > 0.999, + "cosine": cosine, + "max_abs_error": error.max().item(), + "max_scaled_error": scaled_error.max().item(), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kernel", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--mode", choices=("baseline", "sweep", "selected"), default="baseline" + ) + parser.add_argument("--config", type=parse_config) + parser.add_argument("--m", type=int, action="append") + parser.add_argument("--config-shard", type=int, default=0) + parser.add_argument("--num-config-shards", type=int, default=1) + parser.add_argument("--repeats", type=int, default=21) + parser.add_argument("--replays", type=int, default=200) + parser.add_argument("--cache-multiplier", type=float, default=3.0) + parser.add_argument("--max-buffers", type=int, default=32) + parser.add_argument("--max-registers", type=int, default=64) + args = parser.parse_args() + + token_counts = args.m or list(range(1, 17)) + if any(not 1 <= m <= 16 for m in token_counts): + raise ValueError("expected 1 <= M <= 16") + if not 0 <= args.config_shard < args.num_config_shards: + raise ValueError("config shard must be in [0, num_config_shards)") + torch.accelerator.set_device_index(0) + if torch.cuda.get_device_capability() != (10, 3): + raise RuntimeError("this benchmark requires SM103") + + kernel_class = load_kernel_class(args.kernel) + properties = torch.cuda.get_device_properties(0) + metadata = { + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability()), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as output_file: + for m in token_counts: + configs = candidate_configs(args.mode, args.config, m) + torch.manual_seed(20260722 + m) + count = rotating_buffer_count(m, args.cache_multiplier, args.max_buffers) + activation = torch.randn((m, K), device="cuda", dtype=torch.bfloat16) + weights = [ + torch.randn((N, K), device="cuda", dtype=torch.bfloat16) + for _ in range(count) + ] + residuals = [ + torch.randn((m, N), device="cuda", dtype=torch.bfloat16) + for _ in range(count) + ] + candidates: list[tuple[str, Config | None]] = [("cublas_addmm", None)] + candidates.extend( + ("cute_residual", config) + for index, config in enumerate(configs) + if index % args.num_config_shards == args.config_shard + ) + for backend, config in candidates: + row: dict[str, Any] = { + "m": m, + "n": N, + "k": K, + "backend": backend, + "mode": args.mode, + "config": dataclasses.asdict(config) if config else {}, + "num_buffers": count, + "cache_multiplier": args.cache_multiplier, + **metadata, + } + try: + if backend == "cublas_addmm": + launch = lambda a, b, residual, c: torch.addmm( + residual, a, b.t(), out=c + ) + else: + if config is None: + raise AssertionError("missing CuTe config") + compiled = compile_kernel( + kernel_class, m, config, args.max_registers + ) + launch = lambda a, b, residual, c, fn=compiled: fn( + a, b, residual, c, stream() + ) + row.update(resource_usage(compiled)) + samples, outputs = graph_samples( + launch, + activation, + weights, + residuals, + args.repeats, + args.replays, + ) + row.update( + correctness(outputs[0], activation, weights[0], residuals[0]) + ) + row.update(summarize(samples)) + except Exception as error: # noqa: BLE001 + row.update( + { + "valid": False, + "error": f"{type(error).__name__}: {error}", + } + ) + output_file.write(json.dumps(row, sort_keys=True) + "\n") + output_file.flush() + print(json.dumps(row, sort_keys=True), flush=True) + + del activation, weights, residuals + torch.accelerator.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py new file mode 100644 index 00000000000..3c1391919f8 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py @@ -0,0 +1,806 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the Kimi K3 latent-MoE tail and its up-projection kernels. + +The ``up-projection`` subcommand isolates the TP-local dynamic and static-M +skinny GEMMs. It rotates weights through a working set larger than L2 to model +successive model layers. + +The ``whole-tail`` subcommand measures the distributed operator. Its reference +path includes two AllReduces, RMSNorm, the replicated up-projection, and the +final add. CUDA-event samples report the slowest rank so cross-rank skew is +included. + +Examples: + +.. code-block:: console + + .venv/bin/python \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py up-projection + + torchrun --nproc-per-node=8 \ + benchmarks/kernels/benchmark_kimi_k3_latent_moe_tail.py whole-tail + +For multi-node runs, launch one ``torchrun`` agent per node and use a shared +rendezvous endpoint. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +from collections.abc import Callable, Sequence +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import cutlass +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.nn.functional as F +from cuda.bindings import driver as cuda + +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, + set_custom_all_reduce, +) +from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from vllm.models.kimi_k3.nvidia.ops import latent_moe_tail +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail import ( + fused_add_multicast_gemm, + fused_add_multicast_skinny_gemm, +) + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +RMS_EPS = 0.1 +MAX_NUM_TOKENS = 16 +MMA_TILER_MN = (64, 32) +CLUSTER_SHAPE_MN = (1, 8) +B_PRIME_STAGES = 2 + + +def parse_up_projection_config( + value: str, +) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]" + ) from error + if len(values) in (3, 4): + return fused_add_multicast_skinny_gemm.SkinnyConfig(*values) + if len(values) == 5 and values[4] in (0, 1): + return fused_add_multicast_skinny_gemm.SkinnyConfig( + *values[:4], + prefetch_b_before_pdl=bool(values[4]), + ) + raise argparse.ArgumentTypeError( + "config must be BLOCK,OUTPUTS,K_UNROLL" + "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_tail_skinny_config( + value: str, +) -> tuple[int, fused_add_multicast_skinny_gemm.SkinnyConfig]: + try: + values = [int(part) for part in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL[,VECTOR_WIDTH[,PREFETCH_B]]" + ) from error + if len(values) == 4: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 5: + num_tokens, *config = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig(*config) + if len(values) == 6 and values[5] in (0, 1): + num_tokens, block, outputs, unroll, vector_width, prefetch = values + return num_tokens, fused_add_multicast_skinny_gemm.SkinnyConfig( + block, + outputs, + unroll, + vector_width, + bool(prefetch), + ) + raise argparse.ArgumentTypeError( + "config must be M,BLOCK,OUTPUTS,K_UNROLL" + "[,VECTOR_WIDTH[,PREFETCH_B]], where PREFETCH_B is 0 or 1" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="scope", required=True) + + up_projection = subparsers.add_parser( + "up-projection", + help="Benchmark the isolated TP-local up-projection kernels.", + ) + up_projection.add_argument( + "--backend", + choices=("dynamic", "skinny", "both"), + default="both", + ) + up_projection.add_argument("--tp-size", type=int, default=16) + up_projection.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[*range(1, 9), 16], + ) + up_projection.add_argument( + "--skinny-config", + type=parse_up_projection_config, + action="append", + help="Benchmark a static-M config for every selected token count.", + ) + up_projection.add_argument("--cache-multiplier", type=float, default=2.0) + up_projection.add_argument("--max-weights", type=int, default=64) + up_projection.add_argument("--warmup-replays", type=int, default=10) + up_projection.add_argument("--samples", type=int, default=31) + up_projection.add_argument("--output", type=Path) + + whole_tail = subparsers.add_parser( + "whole-tail", + help="Benchmark the distributed latent-MoE tail operator.", + ) + whole_tail.add_argument( + "--backend", + choices=("reference", "fused", "both"), + default="both", + ) + whole_tail.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[1, 5, 8, 16], + ) + whole_tail.add_argument("--warmup-replays", type=int, default=20) + whole_tail.add_argument("--samples", type=int, default=51) + whole_tail.add_argument( + "--skinny-max-num-tokens", + type=int, + nargs="+", + help="Override the fused operator's static-M cutoff; use 0 for dynamic-only.", + ) + whole_tail.add_argument( + "--skinny-config", + type=parse_tail_skinny_config, + action="append", + help="Override one static-M config for tuning.", + ) + whole_tail.add_argument("--output", type=Path) + return parser.parse_args() + + +def percentile(samples: Sequence[float], fraction: float) -> float: + ordered = sorted(samples) + position = fraction * (len(ordered) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + upper_weight = position - lower + return ordered[lower] * (1.0 - upper_weight) + ordered[upper] * upper_weight + + +def summarize(samples_us: Sequence[float]) -> dict[str, Any]: + mean_us = statistics.mean(samples_us) + return { + "median_us": statistics.median(samples_us), + "p10_us": percentile(samples_us, 0.1), + "p90_us": percentile(samples_us, 0.9), + "mean_us": mean_us, + "cv_pct": statistics.pstdev(samples_us) / mean_us * 100.0, + "samples_us": list(samples_us), + } + + +def rotating_weight_count( + shard_size: int, + cache_multiplier: float, + limit: int, +) -> int: + properties = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ) + weight_bytes = shard_size * LATENT_SIZE * 2 + target_bytes = math.ceil(properties.L2_cache_size * cache_multiplier) + return max(2, min(limit, math.ceil(target_bytes / weight_bytes))) + + +def capture_up_projection_graph( + launches: Sequence[Callable[[], None]], +) -> torch.cuda.CUDAGraph: + for launch in launches: + launch() + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for launch in launches: + launch() + torch.accelerator.synchronize() + return graph + + +def benchmark_up_projection_graph( + graph: torch.cuda.CUDAGraph, + *, + operations_per_replay: int, + warmup_replays: int, + samples: int, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + samples_us = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + start.record() + graph.replay() + end.record() + end.synchronize() + samples_us.append(start.elapsed_time(end) * 1000.0 / operations_per_replay) + return summarize(samples_us) + + +class DynamicKernel: + def __init__( + self, + shard_size: int, + mailbox: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + self.shard_size = shard_size + self.mailbox = mailbox + self.mailbox_c = fused_add_multicast_gemm._as_cute(mailbox) + compile_latent = torch.empty( + (1, MAX_NUM_TOKENS, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + compile_weight = torch.empty( + (1, shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=mailbox.device, + ) + cluster_size = math.prod(CLUSTER_SHAPE_MN) + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + self.compiled = fused_add_multicast_gemm.compile_kernel( + (MAX_NUM_TOKENS, shard_size, LATENT_SIZE, 1), + fused_add_multicast_gemm._as_cute( + compile_latent, + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(compile_weight), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + HIDDEN_SIZE, + shard_size, + MMA_TILER_MN, + CLUSTER_SHAPE_MN, + max_active_clusters, + B_PRIME_STAGES, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> None: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + self.compiled( + fused_add_multicast_gemm._as_cute( + latent.unsqueeze(0), + dynamic_m=True, + ), + fused_add_multicast_gemm._as_cute(weight.unsqueeze(0)), + self.mailbox_c, + fused_add_multicast_gemm._as_cute(shared_shard), + cutlass.Int64(latent.shape[0]), + cutlass.Int64(self.mailbox.data_ptr()), + stream, + ) + + +class SkinnyKernel: + def __init__( + self, + num_tokens: int, + shard_size: int, + config: fused_add_multicast_skinny_gemm.SkinnyConfig, + ) -> None: + self.compiled = fused_add_multicast_skinny_gemm.compile_kernel( + num_rows=num_tokens, + latent_dim=LATENT_SIZE, + hidden_dim=HIDDEN_SIZE, + shard_dim=shard_size, + config=config, + ) + + def launch( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + ) -> None: + self.compiled( + fused_add_multicast_skinny_gemm._as_cute(latent), + fused_add_multicast_skinny_gemm._as_cute(weight), + fused_add_multicast_skinny_gemm._as_cute(shared_shard), + cutlass.Int64(mailbox.data_ptr()), + cuda.CUstream(torch.cuda.current_stream().cuda_stream), + ) + + +def check_up_projection_output( + actual: torch.Tensor, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, +) -> None: + gemm = F.linear(latent.float(), weight.float()).to(torch.bfloat16) + expected = (gemm.float() + shared_shard.float()).to(torch.bfloat16) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +def make_up_projection_launches( + launch: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], None], + latent: torch.Tensor, + weights: Sequence[torch.Tensor], + shared_shard: torch.Tensor, +) -> list[Callable[[], None]]: + return [ + lambda weight=weight: launch(latent, weight, shared_shard) for weight in weights + ] + + +def benchmark_up_projection(args: argparse.Namespace) -> None: + if args.tp_size <= 0 or HIDDEN_SIZE % args.tp_size: + raise ValueError("TP size must be positive and divide the hidden size") + if any(not 1 <= num_tokens <= MAX_NUM_TOKENS for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.cache_multiplier <= 0 or args.max_weights <= 0: + raise ValueError("cache multiplier and max weights must be positive") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + + torch.accelerator.set_device_index(0) + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + shard_size = HIDDEN_SIZE // args.tp_size + weight_count = rotating_weight_count( + shard_size, + args.cache_multiplier, + args.max_weights, + ) + torch.manual_seed(20260726) + weights = [ + torch.randn( + (shard_size, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + for _ in range(weight_count) + ] + mailbox = torch.empty( + (1, MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared = torch.randn( + (MAX_NUM_TOKENS, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + shared_shard = shared[:, :shard_size] + use_dynamic = args.backend in ("dynamic", "both") + use_skinny = args.backend in ("skinny", "both") + dynamic_kernel = ( + DynamicKernel(shard_size, mailbox, shared_shard) if use_dynamic else None + ) + + results = [] + for num_tokens in args.num_tokens: + latent = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if dynamic_kernel is not None: + launches = make_up_projection_launches( + dynamic_kernel.launch, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + result["dynamic"] = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + if use_skinny: + configs = args.skinny_config or [ + fused_add_multicast_skinny_gemm.config_for_m( + num_tokens, + shard_size, + ) + ] + skinny_results = [] + for config in configs: + skinny_kernel = SkinnyKernel(num_tokens, shard_size, config) + + def launch_skinny( + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + *, + skinny_kernel: SkinnyKernel = skinny_kernel, + num_tokens: int = num_tokens, + ) -> None: + skinny_kernel.launch( + latent, + weight, + shared_shard[:num_tokens], + mailbox, + ) + + launches = make_up_projection_launches( + launch_skinny, + latent, + weights, + shared_shard, + ) + graph = capture_up_projection_graph(launches) + timing = benchmark_up_projection_graph( + graph, + operations_per_replay=len(launches), + warmup_replays=args.warmup_replays, + samples=args.samples, + ) + check_up_projection_output( + mailbox[0, :num_tokens, :shard_size], + latent, + weights[-1], + shared_shard[:num_tokens], + ) + skinny_results.append( + { + "config": asdict(config), + **timing, + } + ) + result["skinny"] = skinny_results + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "up-projection", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "tp_size": args.tp_size, + "shard_size": shard_size, + "weight_count": weight_count, + "cache_multiplier": args.cache_multiplier, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "results": results, + } + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +def capture_tail_graph( + operation: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]: + for _ in range(3): + dist.barrier(group=cpu_group) + output = operation() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = operation() + torch.accelerator.synchronize() + return graph, output + + +def benchmark_tail_graph( + graph: torch.cuda.CUDAGraph, + *, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, Any]: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + dist.barrier(group=cpu_group) + starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples + 1)] + for start, end in zip(starts, ends): + start.record() + graph.replay() + end.record() + torch.accelerator.synchronize() + + samples_us = torch.tensor( + [start.elapsed_time(end) * 1000.0 for start, end in zip(starts, ends)], + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(samples_us, op=dist.ReduceOp.MAX, group=device_group) + return summarize(samples_us[1:].tolist()) + + +def make_inputs( + num_tokens: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(20260726 + 100 * num_tokens + rank) + routed = torch.randn( + (num_tokens, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ).mul_(0.01) + shared = torch.randn( + (num_tokens, HIDDEN_SIZE), + dtype=torch.bfloat16, + device=device, + ) + return routed, shared + + +def make_reference( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + device_group: dist.ProcessGroup, +) -> Callable[[], torch.Tensor]: + routed_workspace = torch.empty_like(routed) + shared_workspace = torch.empty_like(shared) + + def reference() -> torch.Tensor: + routed_workspace.copy_(routed) + dist.all_reduce(routed_workspace, group=device_group) + normalized = F.rms_norm( + routed_workspace, + (LATENT_SIZE,), + rms_weight, + RMS_EPS, + ) + projected = F.linear(normalized, up_weight) + shared_workspace.copy_(shared) + dist.all_reduce(shared_workspace, group=device_group) + return projected.add(shared_workspace) + + return reference + + +def check_fused_output( + fused_output: torch.Tensor, + reference: Callable[[], torch.Tensor], + cpu_group: dist.ProcessGroup, +) -> None: + dist.barrier(group=cpu_group) + expected = reference() + torch.testing.assert_close(fused_output, expected, atol=8e-2, rtol=3e-2) + + +def benchmark_whole_tail(args: argparse.Namespace) -> None: + if any(not 1 <= num_tokens <= 16 for num_tokens in args.num_tokens): + raise ValueError("--num-tokens values must be in [1, 16]") + if args.warmup_replays < 0 or args.samples <= 0: + raise ValueError("warmup replays must be nonnegative and samples positive") + if args.skinny_max_num_tokens is not None and any( + not 0 <= cutoff <= 8 for cutoff in args.skinny_max_num_tokens + ): + raise ValueError("--skinny-max-num-tokens must be in [0, 8]") + skinny_configs = dict(args.skinny_config or ()) + if len(skinny_configs) != len(args.skinny_config or ()): + raise ValueError("--skinny-config must not repeat an M value") + if any(not 1 <= num_tokens <= 8 for num_tokens in skinny_configs): + raise ValueError("--skinny-config M values must be in [1, 8]") + if not {"RANK", "WORLD_SIZE", "LOCAL_RANK"} <= os.environ.keys(): + raise RuntimeError("launch this benchmark with torchrun") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + device = torch.device("cuda", local_rank) + torch.accelerator.set_device_index(device) + init_distributed_environment() + if world_size > 8: + set_custom_all_reduce(False) + initialize_model_parallel(tensor_model_parallel_size=world_size) + device_group = get_tp_group().device_group + cpu_group = dist.new_group(backend="gloo") + + if torch.cuda.get_device_capability(device)[0] != 10: + raise RuntimeError("Kimi K3 latent-MoE tail requires SM100") + + torch.manual_seed(20260726) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + ) + up_weight = ( + torch.randn( + (HIDDEN_SIZE, LATENT_SIZE), + dtype=torch.bfloat16, + device=device, + ) + / LATENT_SIZE**0.5 + ) + + use_reference = args.backend in ("reference", "both") + use_fused = args.backend in ("fused", "both") + fused_ops = [] + if use_fused: + production_config_for_m = fused_add_multicast_skinny_gemm.config_for_m + + def config_for_m( + num_rows: int, + shard_dim: int = 896, + ) -> fused_add_multicast_skinny_gemm.SkinnyConfig: + config = skinny_configs.get(num_rows) + if config is not None: + return config + return production_config_for_m(num_rows, shard_dim) + + fused_add_multicast_skinny_gemm.config_for_m = config_for_m + cutoffs = args.skinny_max_num_tokens or [latent_moe_tail._SKINNY_MAX_NUM_TOKENS] + for cutoff in cutoffs: + latent_moe_tail._SKINNY_MAX_NUM_TOKENS = cutoff + latent_moe_tail.KimiK3LatentMoETailOp._instances.clear() + fused_ops.append( + ( + cutoff, + latent_moe_tail.KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=RMS_EPS, + ), + ) + ) + cutedsl_warmup() + + results = [] + for num_tokens in args.num_tokens: + routed, shared = make_inputs(num_tokens, rank, device) + reference = make_reference( + routed, + shared, + rms_weight, + up_weight, + device_group, + ) + result: dict[str, Any] = {"num_tokens": num_tokens} + if use_reference: + reference_graph, _ = capture_tail_graph(reference, cpu_group) + result["reference"] = benchmark_tail_graph( + reference_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + for cutoff, fused_op in fused_ops: + + def fused( + routed: torch.Tensor = routed, + shared: torch.Tensor = shared, + fused_op: latent_moe_tail.KimiK3LatentMoETailOp = fused_op, + ) -> torch.Tensor: + return fused_op(routed, shared, rms_weight, up_weight) + + fused_graph, fused_output = capture_tail_graph(fused, cpu_group) + fused_key = "fused" if len(fused_ops) == 1 else f"fused_skinny_max_{cutoff}" + result[fused_key] = benchmark_tail_graph( + fused_graph, + warmup_replays=args.warmup_replays, + samples=args.samples, + device_group=device_group, + cpu_group=cpu_group, + ) + check_fused_output(fused_output, reference, cpu_group) + if "reference" in result: + speedup = ( + result["reference"]["median_us"] / result[fused_key]["median_us"] + ) + if len(fused_ops) == 1: + result["speedup"] = speedup + else: + result[f"{fused_key}_speedup"] = speedup + results.append(result) + + properties = torch.cuda.get_device_properties(device) + report = { + "scope": "whole-tail", + "device": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "world_size": world_size, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "warmup_replays": args.warmup_replays, + "samples": args.samples, + "skinny_max_num_tokens": [cutoff for cutoff, _ in fused_ops], + "skinny_configs": { + str(num_tokens): asdict(config) + for num_tokens, config in skinny_configs.items() + }, + "timing_scope": { + "reference": ( + "two input copies, two AllReduces, RMSNorm, full replicated " + "up-projection GEMM, and final add" + ), + "fused": ( + "routed AllReduce/RMSNorm plus shared ReduceScatter, sharded " + "up-projection/multicast, and Lamport copy" + ), + }, + "results": results, + } + if rank == 0: + rendered = json.dumps(report, indent=2) + print(rendered, flush=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + dist.barrier(group=cpu_group) + + +def main() -> None: + args = parse_args() + if args.scope == "up-projection": + benchmark_up_projection(args) + return + + from vllm.config import VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + benchmark_whole_tail(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py new file mode 100644 index 00000000000..9a244e6c826 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_sp_collectives.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +import json +import os +import statistics +from collections.abc import Callable + +import torch +import torch.distributed as dist + +import vllm._custom_ops as ops +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024]) + parser.add_argument("--hidden-size", type=int, default=7168) + parser.add_argument("--graph-repeats", type=int, default=20) + parser.add_argument("--warmup-replays", type=int, default=5) + parser.add_argument("--samples", type=int, default=15) + return parser.parse_args() + + +def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + op() + stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(repeats): + op() + torch.cuda.current_stream().wait_stream(stream) + return graph + + +def max_rank_graph_time( + graph: torch.cuda.CUDAGraph, + repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> float: + for _ in range(warmup_replays): + graph.replay() + torch.accelerator.synchronize() + + timings = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(samples): + dist.barrier(group=cpu_group) + start.record() + graph.replay() + end.record() + end.synchronize() + elapsed = torch.tensor( + start.elapsed_time(end) / repeats, + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) + timings.append(elapsed.item()) + return statistics.median(timings) + + +def check_outputs( + comm: CustomAllreduce, + local: torch.Tensor, + reduce_input: torch.Tensor, + device_group: dist.ProcessGroup, +) -> None: + expected_gather = torch.empty( + (local.shape[0] * dist.get_world_size(), local.shape[1]), + dtype=local.dtype, + device=local.device, + ) + dist.all_gather_into_tensor(expected_gather, local, group=device_group) + gathered = comm.custom_all_gather(local) + assert gathered is not None + torch.testing.assert_close(gathered, expected_gather) + + expected_scatter = torch.empty_like(local) + dist.reduce_scatter_tensor( + expected_scatter, + reduce_input.clone(), + group=device_group, + ) + scattered = comm.custom_reduce_scatter(reduce_input) + assert scattered is not None + torch.testing.assert_close(scattered, expected_scatter) + + +def benchmark_shape( + comm: CustomAllreduce, + global_tokens: int, + hidden_size: int, + graph_repeats: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, +) -> dict[str, float | int]: + world_size = dist.get_world_size() + rank = dist.get_rank() + padded_tokens = (global_tokens + world_size - 1) // world_size * world_size + local_tokens = padded_tokens // world_size + local = torch.full( + (local_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=torch.accelerator.current_device_index(), + ) + reduce_input = torch.full( + (padded_tokens, hidden_size), + rank + 1, + dtype=torch.bfloat16, + device=local.device, + ) + check_outputs(comm, local, reduce_input, device_group) + + custom_gather_out = torch.empty( + (padded_tokens, hidden_size), + dtype=local.dtype, + device=local.device, + ) + custom_scatter_out = torch.empty_like(local) + nccl_gather_out = torch.empty_like(custom_gather_out) + nccl_scatter_out = torch.empty_like(local) + + def custom_ag() -> None: + ops.mnnvl_lamport_all_gather( + comm._ptr, + local, + custom_gather_out, + comm.mnnvl_lamport_ag_local_ptr, + comm.mnnvl_lamport_ag_multicast_ptr, + comm.mnnvl_lamport_ag_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def custom_rs() -> None: + ops.mnnvl_lamport_reduce_scatter( + comm._ptr, + reduce_input, + custom_scatter_out, + comm.mnnvl_lamport_rs_local_ptr, + comm.mnnvl_lamport_rs_epoch_ptr, + comm.mnnvl_buffer_size, + ) + + def nccl_ag() -> None: + dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group) + + def nccl_rs() -> None: + dist.reduce_scatter_tensor( + nccl_scatter_out, + reduce_input, + group=device_group, + ) + + graphs = { + "custom_ag_us": capture_graph(custom_ag, graph_repeats), + "nccl_ag_us": capture_graph(nccl_ag, graph_repeats), + "custom_rs_us": capture_graph(custom_rs, graph_repeats), + "nccl_rs_us": capture_graph(nccl_rs, graph_repeats), + } + times = { + name: max_rank_graph_time( + graph, + graph_repeats, + warmup_replays, + samples, + device_group, + cpu_group, + ) + * 1000 + for name, graph in graphs.items() + } + torch.testing.assert_close(custom_gather_out, nccl_gather_out) + torch.testing.assert_close(custom_scatter_out, nccl_scatter_out) + return { + "global_tokens": global_tokens, + "padded_tokens": padded_tokens, + "local_bytes": local.nbytes, + "full_bytes": reduce_input.nbytes, + **times, + "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"], + "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"], + } + + +def main() -> None: + args = parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + dist.init_process_group("nccl") + device_group = dist.group.WORLD + cpu_group = dist.new_group(backend="gloo") + + comm = CustomAllreduce( + group=cpu_group, + device=torch.device("cuda", local_rank), + ) + assert not comm.disabled + assert comm.world_size == 16 + assert comm.mnnvl_only + assert comm.mnnvl_multicast_ptr + + results = [ + benchmark_shape( + comm, + tokens, + args.hidden_size, + args.graph_repeats, + args.warmup_replays, + args.samples, + device_group, + cpu_group, + ) + for tokens in args.tokens + ] + if dist.get_rank() == 0: + print(json.dumps(results, indent=2), flush=True) + + comm.close() + dist.destroy_process_group(cpu_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake new file mode 100644 index 00000000000..28b471a9e74 --- /dev/null +++ b/cmake/external_projects/flashkda.cmake @@ -0,0 +1,74 @@ +include(FetchContent) + +if(DEFINED ENV{FLASH_KDA_SRC_DIR}) + set(FLASH_KDA_SRC_DIR $ENV{FLASH_KDA_SRC_DIR}) +endif() + +if(FLASH_KDA_SRC_DIR) + FetchContent_Declare( + flashkda + SOURCE_DIR ${FLASH_KDA_SRC_DIR} + ) +else() + FetchContent_Declare( + flashkda + GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git + GIT_TAG a3e42bbbece3bb38f7c426b880315294a336e82f + GIT_PROGRESS TRUE + GIT_SUBMODULES cutlass + ) +endif() + +FetchContent_MakeAvailable(flashkda) +message(STATUS "FlashKDA is available at ${flashkda_SOURCE_DIR}") + +set(FLASH_KDA_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "9.0a") +endif() +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0f" "12.0f") +elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND FLASH_KDA_SUPPORT_ARCHS "10.0a" "10.3a" "12.0a") +endif() + +cuda_archs_loose_intersection( + FLASH_KDA_ARCHS "${FLASH_KDA_SUPPORT_ARCHS}" "${CUDA_ARCHS}") + +if(FLASH_KDA_ARCHS) + message(STATUS "FlashKDA CUDA architectures: ${FLASH_KDA_ARCHS}") + + set(FLASH_KDA_SOURCES + csrc/flashkda_registration.cpp + ${flashkda_SOURCE_DIR}/csrc/flash_kda.cpp + ${flashkda_SOURCE_DIR}/csrc/smxx/fwd_launch.cu) + set(FLASH_KDA_INCLUDES + ${flashkda_SOURCE_DIR}/csrc + ${flashkda_SOURCE_DIR}/cutlass/include + ${flashkda_SOURCE_DIR}/cutlass/examples/common + ${flashkda_SOURCE_DIR}/cutlass/tools/util/include) + + set_gencode_flags_for_srcs( + SRCS "${FLASH_KDA_SOURCES}" + CUDA_ARCHS "${FLASH_KDA_ARCHS}") + + define_extension_target( + _flashkda_C + DESTINATION vllm + LANGUAGE ${VLLM_GPU_LANG} + SOURCES ${FLASH_KDA_SOURCES} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES ${FLASH_KDA_INCLUDES} + USE_SABI 3 + WITH_SOABI) + + target_compile_options(_flashkda_C PRIVATE + $<$:-UPy_LIMITED_API --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -O3> + $<$:-UPy_LIMITED_API>) +else() + message(STATUS + "FlashKDA will not compile: CUDA >=12.0 and a supported architecture " + "(SM90, SM10x, or SM12x) are required") + add_custom_target(_flashkda_C) +endif() diff --git a/csrc/custom_all_gather_reduce_scatter.cuh b/csrc/custom_all_gather_reduce_scatter.cuh new file mode 100644 index 00000000000..ac990d8106a --- /dev/null +++ b/csrc/custom_all_gather_reduce_scatter.cuh @@ -0,0 +1,327 @@ +#pragma once + +#include "custom_collective_common.cuh" + +namespace vllm { + +constexpr int kMnnvlLamportAgThreads = 128; +constexpr int kMnnvlLamportRsThreads = 256; +constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192; + +using CopyPack = array_t; + +template +__global__ void __launch_bounds__(512, 1) + cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg, + CopyPack* __restrict__ result, int rank, + int size_per_rank) { + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + barrier_at_start(sg, self_sg, rank); +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + auto src = reinterpret_cast(dp.ptrs[src_rank]); + auto dst = result + src_rank * size_per_rank; + for (int idx = tid; idx < size_per_rank; idx += stride) { + dst[idx] = src[idx]; + } + } + barrier_at_end(sg, self_sg, rank); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, + int size_per_rank) { + using P = typename packed_t::P; + using A = typename packed_t::A; + auto dp = *_dp; + auto offset = rank * size_per_rank; + barrier_at_start(sg, self_sg, rank); + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank; + idx += gridDim.x * blockDim.x) { + reinterpret_cast(result)[idx] = + packed_reduce((const P**)&dp.ptrs[0], offset + idx); + } + barrier_at_end(sg, self_sg, rank); +} + +template +union LamportPack { + P packed; + uint32_t words[sizeof(P) / sizeof(uint32_t)]; +}; + +template +DINLINE LamportPack

load_lamport_pack(const P* ptr) { + static_assert(sizeof(P) == 16); + LamportPack

value; +#if !defined(USE_ROCM) + asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(value.words[0]), "=r"(value.words[1]), + "=r"(value.words[2]), "=r"(value.words[3]) + : "l"(ptr) + : "memory"); +#else + const volatile uint32_t* src = + reinterpret_cast(ptr); + #pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = src[i]; + } +#endif + return value; +} + +template +DINLINE bool is_lamport_dirty(const LamportPack

& value) { +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) return true; + } + return false; +} + +template +DINLINE P lamport_sentinel() { + LamportPack

value; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + value.words[i] = 0x80000000U; + } + return value.packed; +} + +template +DINLINE P sanitize_lamport_payload(P packed) { + LamportPack

value{.packed = packed}; +#pragma unroll + for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { + if (value.words[i] == 0x80000000U) value.words[i] = 0; + } + return value.packed; +} + +template +DINLINE P wait_lamport_payload(const P* ptr) { + auto value = load_lamport_pack(ptr); + while (is_lamport_dirty(value)) value = load_lamport_pack(ptr); + return value.packed; +} + +template +DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride, + P local_value, P (&values)[ngpus]) { + bool ready[ngpus]; +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + ready[src_rank] = src_rank == rank; + if (src_rank == rank) values[src_rank] = local_value; + } + + int remaining = ngpus - 1; + while (remaining != 0) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + if (!ready[src_rank]) { + auto value = load_lamport_pack(base + src_rank * rank_stride); + if (!is_lamport_dirty(value)) { + values[src_rank] = value.packed; + ready[src_rank] = true; + --remaining; + } + } + } + } +} + +template +DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input, + int rank, int size_per_rank, int idx) { + P source_zero = + rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx); + A tmp = upcast(source_zero); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + P value = src_rank == rank + ? packed_input[rank * size_per_rank + idx] + : wait_lamport_payload(current_local + + src_rank * size_per_rank + idx); + packed_assign_add(tmp, upcast(value)); + } + return sanitize_lamport_payload(downcast

(tmp)); +} + +DINLINE void lamport_cta_arrive(uint32_t* counter) { +#if !defined(USE_ROCM) + if (threadIdx.x < 32) { + asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory"); + if (threadIdx.x == 0) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("red.release.global.gpu.add.u32 [%0], 1;" + : + : "l"(counter) + : "memory"); + #else + atomicAdd(counter, 1); + #endif + } + } else { + asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory"); + } +#else + __syncthreads(); + if (threadIdx.x == 0) atomicAdd(counter, 1); +#endif +} + +template +__global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1) + mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input, + T* __restrict__ result, + T* __restrict__ multicast_buffer, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto current_multicast = + reinterpret_cast(multicast_buffer) + current_stage * stage_size; + auto packed_input = reinterpret_cast(input); + auto packed_result = reinterpret_cast(result); + + int total_size = size_per_rank * ngpus; + P local_value; + if (tid < size_per_rank) { + local_value = packed_input[tid]; + current_multicast[rank * size_per_rank + tid] = + sanitize_lamport_payload(local_value); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (tid < size_per_rank) { +#pragma unroll + for (int src_rank = 0; src_rank < ngpus; ++src_rank) { + int output_idx = src_rank * size_per_rank + tid; + P value = src_rank == rank + ? local_value + : wait_lamport_payload(current_local + output_idx); + packed_result[output_idx] = value; + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = total_size; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +template +__global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1) + mnnvl_lamport_reduce_scatter_kernel(RankData* _dp, + const T* __restrict__ input, + T* __restrict__ result, + uint32_t* __restrict__ epochs, int rank, + int size_per_rank, int stage_size) { + using P = typename packed_t::P; + using A = typename packed_t::A; +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + auto dp = *_dp; + int dst_rank = blockIdx.x % ngpus; + int tile = blockIdx.x / ngpus; + int idx = tile * blockDim.x + threadIdx.x; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + uint32_t epoch = epochs[0]; + int current_stage = epoch % 3; + int dirty_stage = (epoch + 1) % 3; + int dirty_size = epochs[2 + dirty_stage]; + auto local_buffer = reinterpret_cast(const_cast(dp.ptrs[rank])); + auto current_local = local_buffer + current_stage * stage_size; + auto dirty_local = local_buffer + dirty_stage * stage_size; + auto packed_input = reinterpret_cast(input); + + if (idx < size_per_rank && dst_rank != rank) { + auto dst = reinterpret_cast(const_cast(dp.ptrs[dst_rank])) + + current_stage * stage_size + rank * size_per_rank; + auto src = packed_input + dst_rank * size_per_rank; + dst[idx] = sanitize_lamport_payload(src[idx]); + } +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + lamport_cta_arrive(&epochs[1]); + + for (int idx = tid; idx < dirty_size; idx += stride) { + dirty_local[idx] = lamport_sentinel

(); + } + + if (idx < size_per_rank && dst_rank == rank) { + if constexpr (ngpus == 4) { + if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) { + reinterpret_cast(result)[idx] = + reduce_lamport_payloads(current_local, packed_input, + rank, size_per_rank, idx); + } else { + P values[ngpus]; + wait_lamport_payloads( + current_local + idx, rank, size_per_rank, + packed_input[rank * size_per_rank + idx], values); + A tmp = upcast(values[0]); +#pragma unroll + for (int src_rank = 1; src_rank < ngpus; ++src_rank) { + packed_assign_add(tmp, upcast(values[src_rank])); + } + reinterpret_cast(result)[idx] = + sanitize_lamport_payload(downcast

(tmp)); + } + } else { + reinterpret_cast(result)[idx] = reduce_lamport_payloads( + current_local, packed_input, rank, size_per_rank, idx); + } + } + + if (tid == 0) { + while (*reinterpret_cast(&epochs[1]) < gridDim.x); + epochs[2 + current_stage] = size_per_rank * ngpus; + epochs[0] = epoch + 1; + epochs[1] = 0; + } +} + +} // namespace vllm diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 58926f6429d..385a0e0ae79 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -1,299 +1,8 @@ #pragma once -#include -#include -#include -#include - -#if defined(USE_ROCM) -typedef __hip_bfloat16 nv_bfloat16; -#endif - -#include -#include -#include -#include -#include -#include -#include -#include +#include "custom_collective_common.cuh" namespace vllm { -#define CUDACHECK(cmd) \ - do { \ - cudaError_t e = cmd; \ - if (e != cudaSuccess) { \ - printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ - cudaGetErrorString(e)); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) - -// Maximal number of blocks in allreduce kernel. -constexpr int kMaxBlocks = 36; - -// Default number of blocks in allreduce kernel. -#ifndef USE_ROCM -const int defaultBlockLimit = 36; -CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#else -const int defaultBlockLimit = 16; -hipPointer_attribute rangeStartAddrAttr = - HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; -#endif - -// Counter may overflow, but it's fine since unsigned int overflow is -// well-defined behavior. -using FlagType = uint32_t; - -// Two sets of peer counters are needed for two syncs: starting and ending an -// operation. The reason is that it's possible for peer GPU block to arrive at -// the second sync point while the current GPU block haven't passed the first -// sync point. Thus, peer GPU may write counter+1 while current GPU is busy -// waiting for counter. We use alternating counter array to avoid this -// possibility. -struct Signal { - alignas(128) FlagType start[kMaxBlocks][8]; - alignas(128) FlagType end[kMaxBlocks][8]; - alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank -}; - -struct __align__(16) RankData { - const void* ptrs[8]; -}; - -struct __align__(16) RankSignals { - Signal* signals[8]; -}; - -// like std::array, but aligned -template -struct __align__(alignof(T) * sz) array_t { - T data[sz]; - using type = T; - static constexpr int size = sz; -}; - -// use packed type to maximize memory efficiency -// goal: generate ld.128 and st.128 instructions -template -struct packed_t { - // the (P)acked type for load/store - using P = array_t; - // the (A)ccumulator type for reduction - using A = array_t; -}; - -#define DINLINE __device__ __forceinline__ - -// scalar cast functions -DINLINE float upcast_s(half val) { return __half2float(val); } - -template -DINLINE T downcast_s(float val); -template <> -DINLINE half downcast_s(float val) { - return __float2half(val); -} - -// scalar add functions -// for some reason when compiling with Pytorch, the + operator for half and -// bfloat is disabled so we call the intrinsics directly -DINLINE half& assign_add(half& a, half b) { - a = __hadd(a, b); - return a; -} -DINLINE float& assign_add(float& a, float b) { return a += b; } - -#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) -DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } -template <> -DINLINE nv_bfloat16 downcast_s(float val) { - return __float2bfloat16(val); -} -DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { - a = __hadd(a, b); - return a; -} -#endif - -template -DINLINE array_t& packed_assign_add(array_t& a, array_t b) { -#pragma unroll - for (int i = 0; i < N; i++) { - assign_add(a.data[i], b.data[i]); - } - return a; -} - -template -DINLINE array_t upcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - array_t out; -#pragma unroll - for (int i = 0; i < N; i++) { - out.data[i] = upcast_s(val.data[i]); - } - return out; - } -} - -template -DINLINE O downcast(array_t val) { - if constexpr (std::is_same::value) { - return val; - } else { - O out; -#pragma unroll - for (int i = 0; i < O::size; i++) { - out.data[i] = downcast_s(val.data[i]); - } - return out; - } -} - -#if !defined(USE_ROCM) - -static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #else - asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), - "l"(flag_addr)); - #endif -} - -static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { - FlagType flag; - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 - asm volatile("ld.acquire.sys.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - #else - asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" - : "=r"(flag) - : "l"(flag_addr)); - #endif - return flag; -} - -static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { - asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); -} - -static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { - FlagType flag; - asm volatile("ld.volatile.global.u32 %0, [%1];" - : "=r"(flag) - : "l"(flag_addr)); - return flag; -} - -// This function is meant to be used as the first synchronization in the all -// reduce kernel. Thus, it doesn't need to make any visibility guarantees for -// prior memory accesses. Note: volatile writes will not be reordered against -// other volatile writes. -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value - // from peer. - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -// This function is meant to be used as the second or the final -// synchronization barrier in the all reduce kernel. If it's the final -// synchronization barrier, we don't need to make any visibility guarantees -// for prior memory accesses. -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; - auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; - // Write the expected counter value to peer and wait for correct value from - // peer. - if constexpr (!final_sync) { - st_flag_release(peer_counter_ptr, flag); - while (ld_flag_acquire(self_counter_ptr) != flag); - } else { - st_flag_volatile(peer_counter_ptr, flag); - while (ld_flag_volatile(self_counter_ptr) != flag); - } - } - if constexpr (!final_sync) __syncthreads(); - - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#else - -template -DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, - int rank) { - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], - flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], - __ATOMIC_RELAXED, - __MEMORY_SCOPE_DEVICE) < flag); - } - __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -template -DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { - __syncthreads(); - uint32_t flag = self_sg->_flag[blockIdx.x] + 1; - if (threadIdx.x < ngpus) { - // simultaneously write to the corresponding flag of all ranks. - // Latency = 1 p2p write - __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], - flag, - final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, - __MEMORY_SCOPE_SYSTEM); - // wait until we got true from all ranks - while ( - __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], - final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, - __MEMORY_SCOPE_DEVICE) < flag); - } - if constexpr (!final_sync) __syncthreads(); - // use one thread to update flag - if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; -} - -#endif - -template -DINLINE P packed_reduce(const P* ptrs[], int idx) { - A tmp = upcast(ptrs[0][idx]); -#pragma unroll - for (int i = 1; i < ngpus; i++) { - packed_assign_add(tmp, upcast(ptrs[i][idx])); - } - return downcast

(tmp); -} template __global__ void __launch_bounds__(512, 1) @@ -616,6 +325,21 @@ class CustomAllreduce { #undef KL } + void allgather(cudaStream_t stream, void* input, void* output, int size_bytes, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_allgather(cudaStream_t stream, T* input, T* output, + void* local_buffer, void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes); + template + void reduce_scatter(cudaStream_t stream, T* input, T* output, int size, + int threads = 512, int block_limit = defaultBlockLimit); + template + void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output, + void* local_buffer, uint32_t* epochs, + int size, int stage_size_bytes); + ~CustomAllreduce() { for (auto [_, ptr] : ipc_handles_) { CUDACHECK(cudaIpcCloseMemHandle(ptr)); @@ -625,8 +349,8 @@ class CustomAllreduce { /** * To inspect PTX/SASS, copy paste this header file to compiler explorer and - add a template instantiation: + * add a template instantiation: * template void vllm::CustomAllreduce::allreduce(cudaStream_t, half *, - half *, int, int, int); -*/ -} // namespace vllm \ No newline at end of file + * half *, int, int, int); + */ +} // namespace vllm diff --git a/csrc/custom_collective_common.cuh b/csrc/custom_collective_common.cuh new file mode 100644 index 00000000000..353dcd07e8b --- /dev/null +++ b/csrc/custom_collective_common.cuh @@ -0,0 +1,332 @@ +#pragma once + +#include +#include +#include +#include + +#if defined(USE_ROCM) +typedef __hip_bfloat16 nv_bfloat16; +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vllm { +constexpr int kMaxCustomCollectiveRanks = 16; + +#define CUDACHECK(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \ + cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +// Maximal number of blocks in allreduce kernel. +constexpr int kMaxBlocks = 36; + +// Default number of blocks in allreduce kernel. +#ifndef USE_ROCM +inline constexpr int defaultBlockLimit = 36; +inline CUpointer_attribute rangeStartAddrAttr = + CU_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#else +inline constexpr int defaultBlockLimit = 16; +inline hipPointer_attribute rangeStartAddrAttr = + HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR; +#endif + +// Counter may overflow, but it's fine since unsigned int overflow is +// well-defined behavior. +using FlagType = uint32_t; + +// Two sets of peer counters are needed for two syncs: starting and ending an +// operation. The reason is that it's possible for peer GPU block to arrive at +// the second sync point while the current GPU block haven't passed the first +// sync point. Thus, peer GPU may write counter+1 while current GPU is busy +// waiting for counter. We use alternating counter array to avoid this +// possibility. +struct Signal { + alignas(128) FlagType start[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType end[kMaxBlocks][kMaxCustomCollectiveRanks]; + alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank +}; + +struct __align__(16) RankData { + const void* ptrs[kMaxCustomCollectiveRanks]; +}; + +struct __align__(16) RankSignals { + Signal* signals[kMaxCustomCollectiveRanks]; +}; + +// like std::array, but aligned +template +struct __align__(alignof(T) * sz) array_t { + T data[sz]; + using type = T; + static constexpr int size = sz; +}; + +// use packed type to maximize memory efficiency +// goal: generate ld.128 and st.128 instructions +template +struct packed_t { + // the (P)acked type for load/store + using P = array_t; + // the (A)ccumulator type for reduction + using A = array_t; +}; + +#define DINLINE __device__ __forceinline__ + +// scalar cast functions +DINLINE float upcast_s(half val) { return __half2float(val); } + +template +DINLINE T downcast_s(float val); +template <> +DINLINE half downcast_s(float val) { + return __float2half(val); +} + +// scalar add functions +// for some reason when compiling with Pytorch, the + operator for half and +// bfloat is disabled so we call the intrinsics directly +DINLINE half& assign_add(half& a, half b) { + a = __hadd(a, b); + return a; +} +DINLINE float& assign_add(float& a, float b) { return a += b; } + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } +template <> +DINLINE nv_bfloat16 downcast_s(float val) { + return __float2bfloat16(val); +} +DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { + a = __hadd(a, b); + return a; +} +#endif + +template +DINLINE array_t& packed_assign_add(array_t& a, array_t b) { +#pragma unroll + for (int i = 0; i < N; i++) { + assign_add(a.data[i], b.data[i]); + } + return a; +} + +template +DINLINE array_t upcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + array_t out; +#pragma unroll + for (int i = 0; i < N; i++) { + out.data[i] = upcast_s(val.data[i]); + } + return out; + } +} + +template +DINLINE O downcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + O out; +#pragma unroll + for (int i = 0; i < O::size; i++) { + out.data[i] = downcast_s(val.data[i]); + } + return out; + } +} + +#if !defined(USE_ROCM) + +static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #else + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), + "l"(flag_addr)); + #endif +} + +static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { + FlagType flag; + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + #else + asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" + : "=r"(flag) + : "l"(flag_addr)); + #endif + return flag; +} + +static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { + asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +} + +static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { + FlagType flag; + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(flag) + : "l"(flag_addr)); + return flag; +} + +// This function is meant to be used as the first synchronization in the all +// reduce kernel. Thus, it doesn't need to make any visibility guarantees for +// prior memory accesses. Note: volatile writes will not be reordered against +// other volatile writes. +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value + // from peer. + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x]; + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +// This function is meant to be used as the second or the final +// synchronization barrier in the all reduce kernel. If it's the final +// synchronization barrier, we don't need to make any visibility guarantees +// for prior memory accesses. +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x]; + // Write the expected counter value to peer and wait for correct value from + // peer. + if constexpr (!final_sync) { + st_flag_release(peer_counter_ptr, flag); + while (ld_flag_acquire(self_counter_ptr) != flag); + } else { + st_flag_volatile(peer_counter_ptr, flag); + while (ld_flag_volatile(self_counter_ptr) != flag); + } + } + if constexpr (!final_sync) __syncthreads(); + + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#else + +template +DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg, + int rank) { + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_RELAXED, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_start_release(const RankSignals& sg, Signal* self_sg, + int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank], + flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x], + __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + __syncthreads(); + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +template +DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) { + __syncthreads(); + uint32_t flag = self_sg->_flag[blockIdx.x] + 1; + if (threadIdx.x < ngpus) { + // simultaneously write to the corresponding flag of all ranks. + // Latency = 1 p2p write + __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank], + flag, + final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE, + __MEMORY_SCOPE_SYSTEM); + // wait until we got true from all ranks + while ( + __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x], + final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE, + __MEMORY_SCOPE_DEVICE) < flag); + } + if constexpr (!final_sync) __syncthreads(); + // use one thread to update flag + if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag; +} + +#endif + +template +DINLINE P packed_reduce(const P* ptrs[], int idx) { + A tmp = upcast(ptrs[0][idx]); +#pragma unroll + for (int i = 1; i < ngpus; i++) { + packed_assign_add(tmp, upcast(ptrs[i][idx])); + } + return downcast

(tmp); +} + +} // namespace vllm diff --git a/csrc/flashkda_registration.cpp b/csrc/flashkda_registration.cpp new file mode 100644 index 00000000000..9058119e119 --- /dev/null +++ b/csrc/flashkda_registration.cpp @@ -0,0 +1,17 @@ +#include "core/registration.h" +#include "flash_kda.h" + +TORCH_LIBRARY(_flashkda_C, m) { + m.def("get_workspace_size(int T_total, int H, int N=1) -> int", + &get_workspace_size); + m.def( + "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, " + "Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, " + "float lower_bound, " + "Tensor? initial_state=None, Tensor(b!)? final_state=None, " + "Tensor? cu_seqlens=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); } + +REGISTER_EXTENSION(_flashkda_C) diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index 60b8ca5f382..e76250624d6 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -464,6 +464,66 @@ __global__ void swigluoai_and_mul_kernel( } } +// SITU (Kimi SituGLU) gated activation. Non-interleaved layout: +// input = [gate(d), up(d)] per token. +// gate_out = beta * tanh(gate / beta) * sigmoid(gate) +// up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up +// out = gate_out * up_out +// Compute is done in fp32 and written straight to `out` -- no intermediate +// tensors and no full-tensor fp32 upcast (the pure-torch forward_native +// allocated ~8 fp32 temporaries per call, which blows up MoE profiling). +template +__global__ void situ_and_mul_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., 2, d] + const int d, const float beta, const float linear_beta) { + const int64_t row = blockIdx.x; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + const float g = (float)VLLM_LDG(&gate_ptr[idx]); + const float u = (float)VLLM_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + +template +__global__ void masked_situ_and_mul_kernel( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, + const int* __restrict__ expert_num_tokens, const int max_num_tokens, + const int d, const float beta, const float linear_beta) { + const int expert = blockIdx.y; + const int num_tokens = expert_num_tokens[expert]; + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= d || num_tokens == 0) { + return; + } + + const bool clamp_up = linear_beta > 0.0f; + const float inv_beta = 1.0f / beta; + const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f; + const int64_t expert_row = static_cast(expert) * max_num_tokens; + for (int token = 0; token < num_tokens; ++token) { + const int64_t row = expert_row + token; + const scalar_t* gate_ptr = input + row * 2 * d; + const scalar_t* up_ptr = gate_ptr + d; + scalar_t* out_ptr = out + row * d; + const float g = (float)VLLM_LDG(&gate_ptr[idx]); + const float u = (float)VLLM_LDG(&up_ptr[idx]); + const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g)); + const float up_out = + clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u; + out_ptr[idx] = (scalar_t)(gate_out * up_out); + } +} + } // namespace vllm #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \ @@ -553,6 +613,54 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d] double alpha, double limit) { LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit); } + +// Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed +// through), matching SituAndMul(linear_beta=None) on the Python side. +void situ_and_mul(torch::stable::Tensor& out, // [..., d] + torch::stable::Tensor& input, // [..., 2 * d] + double beta, double linear_beta) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + if (num_tokens == 0) { + return; + } + dim3 grid(num_tokens); + dim3 block(std::min(d, 1024)); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "situ_and_mul_kernel", [&] { + vllm::situ_and_mul_kernel<<>>( + out.mutable_data_ptr(), input.const_data_ptr(), + d, (float)beta, (float)linear_beta); + }); +} + +void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d] + torch::stable::Tensor& input, // [E, T, 2 * d] + const torch::stable::Tensor& expert_num_tokens, + double beta, double linear_beta) { + int num_experts = input.size(0); + int max_num_tokens = input.size(1); + int d = input.size(2) / 2; + if (num_experts == 0 || max_num_tokens == 0) { + return; + } + constexpr int block_size = 256; + dim3 grid((d + block_size - 1) / block_size, num_experts); + dim3 block(block_size); + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "masked_situ_and_mul_kernel", [&] { + vllm::masked_situ_and_mul_kernel<<>>( + out.mutable_data_ptr(), input.const_data_ptr(), + expert_num_tokens.const_data_ptr(), max_num_tokens, d, + (float)beta, (float)linear_beta); + }); +} namespace vllm { // Element-wise activation kernel template. diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index b132e82e253..cc89397f68a 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -21,7 +21,10 @@ __global__ void merge_attn_states_kernel( const float* prefix_lse, const scalar_t* suffix_output, const float* suffix_lse, const uint num_tokens, const uint num_heads, const uint head_size, const uint prefix_head_stride, - const uint output_head_stride, const uint prefix_num_tokens, + const uint output_head_stride, const uint prefix_lse_head_stride, + const uint prefix_lse_token_stride, const uint suffix_lse_head_stride, + const uint suffix_lse_token_stride, const uint output_lse_head_stride, + const uint output_lse_token_stride, const uint prefix_num_tokens, const float* output_scale) { // Inputs always load 128-bit packs (pack_size elements of scalar_t). // Outputs store pack_size elements of output_t, which is smaller for FP8. @@ -84,15 +87,19 @@ __global__ void merge_attn_states_kernel( } } if (output_lse != nullptr && pack_idx == 0) { - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; - output_lse[head_idx * num_tokens + token_idx] = s_lse; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = s_lse; } return; } // For tokens within prefix range, merge prefix and suffix - float p_lse = prefix_lse[head_idx * num_tokens + token_idx]; - float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; + float p_lse = prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; @@ -132,7 +139,8 @@ __global__ void merge_attn_states_kernel( } // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { - output_lse[head_idx * num_tokens + token_idx] = max_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = max_lse; } return; } @@ -187,7 +195,8 @@ __global__ void merge_attn_states_kernel( // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { float out_lse = logf(out_se) + max_lse; - output_lse[head_idx * num_tokens + token_idx] = out_lse; + output_lse[head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride] = out_lse; } } @@ -221,6 +230,9 @@ __global__ void merge_attn_states_kernel( reinterpret_cast(suffix_output.data_ptr()), \ reinterpret_cast(suffix_lse.data_ptr()), num_tokens, \ num_heads, head_size, prefix_head_stride, output_head_stride, \ + prefix_lse_head_stride, prefix_lse_token_stride, \ + suffix_lse_head_stride, suffix_lse_token_stride, \ + output_lse_head_stride, output_lse_token_stride, \ prefix_num_tokens, output_scale_ptr); \ } @@ -259,6 +271,19 @@ void merge_attn_states_launcher( const uint head_size = output.size(2); const uint prefix_head_stride = prefix_output.stride(1); const uint output_head_stride = output.stride(1); + // lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views + // (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index + // them by their actual strides rather than assuming a contiguous layout. + const uint prefix_lse_head_stride = prefix_lse.stride(0); + const uint prefix_lse_token_stride = prefix_lse.stride(1); + const uint suffix_lse_head_stride = suffix_lse.stride(0); + const uint suffix_lse_token_stride = suffix_lse.stride(1); + uint output_lse_head_stride = 0; + uint output_lse_token_stride = 0; + if (output_lse.has_value()) { + output_lse_head_stride = output_lse.value().stride(0); + output_lse_token_stride = output_lse.value().stride(1); + } // Thread mapping is based on input BF16 pack_size const uint pack_size = 16 / sizeof(scalar_t); STD_TORCH_CHECK(head_size % pack_size == 0, diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 5c1628537aa..c35a516b5c9 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -443,6 +443,55 @@ __global__ void concat_and_cache_mla_kernel( copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank); } +// Grouped variant of concat_and_cache_mla: inserts the context K/V for every +// draft layer in a single launch. Grid is (num_tokens, num_layers); each layer +// reads its own cache base pointer from kv_cache_ptrs (same pointer-array +// pattern as copy_blocks_kernel). bf16 only, so it is a raw 16-bit copy with no +// scaling or quantization; scalar_t is uint16_t for portability. +template +__global__ void concat_and_cache_mla_grouped_kernel( + const scalar_t* __restrict__ kv_c, // [num_layers, num_tokens, + // kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_layers, num_tokens, pe_dim] + const int64_t* __restrict__ kv_cache_ptrs, // [num_layers] + const int64_t* __restrict__ slot_mapping, // [num_layers, num_tokens] + const int64_t kv_c_layer_stride, const int64_t kv_c_token_stride, + const int64_t k_pe_layer_stride, const int64_t k_pe_token_stride, + const int64_t slot_layer_stride, const int64_t block_stride, + const int64_t entry_stride, const int kv_lora_rank, const int pe_dim, + const int block_size) { + const int64_t token_idx = blockIdx.x; + const int64_t layer_idx = blockIdx.y; + const int64_t slot_idx = + slot_mapping[layer_idx * slot_layer_stride + token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + scalar_t* __restrict__ kv_cache = + reinterpret_cast(kv_cache_ptrs[layer_idx]); + const scalar_t* __restrict__ kv_c_layer = + kv_c + layer_idx * kv_c_layer_stride; + const scalar_t* __restrict__ k_pe_layer = + k_pe + layer_idx * k_pe_layer_stride; + + auto copy = [&](const scalar_t* __restrict__ src, int64_t src_token_stride, + int size, int offset) { + for (int i = threadIdx.x; i < size; i += blockDim.x) { + const int64_t src_idx = token_idx * src_token_stride + i; + const int64_t dst_idx = + block_idx * block_stride + block_offset * entry_stride + i + offset; + kv_cache[dst_idx] = src[src_idx]; + } + }; + + copy(kv_c_layer, kv_c_token_stride, kv_lora_rank, 0); + copy(k_pe_layer, k_pe_token_stride, pe_dim, kv_lora_rank); +} + template __global__ void concat_and_cache_ds_mla_kernel( const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] @@ -902,6 +951,53 @@ void concat_and_cache_mla( } } +void concat_and_cache_mla_grouped( + torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim] + torch::stable::Tensor& kv_cache_ptrs, // [num_layers] int64, on device + torch::stable::Tensor& slot_mapping, // [num_layers, num_tokens] int64 + int64_t block_size, int64_t block_stride, int64_t entry_stride) { + int num_layers = kv_c.size(0); + int num_tokens = kv_c.size(1); + int kv_lora_rank = kv_c.size(2); + int pe_dim = k_pe.size(2); + + STD_TORCH_CHECK( + kv_c.scalar_type() == torch::headeronly::ScalarType::BFloat16 && + k_pe.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "concat_and_cache_mla_grouped only supports a bf16 KV cache; got kv_c=", + kv_c.scalar_type(), ", k_pe=", k_pe.scalar_type()); + STD_TORCH_CHECK( + kv_cache_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "kv_cache_ptrs must be int64"); + + if (num_tokens == 0 || num_layers == 0) { + return; + } + + const int64_t kv_c_layer_stride = kv_c.stride(0); + const int64_t kv_c_token_stride = kv_c.stride(1); + const int64_t k_pe_layer_stride = k_pe.stride(0); + const int64_t k_pe_token_stride = k_pe.stride(1); + const int64_t slot_layer_stride = slot_mapping.stride(0); + + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + + dim3 grid(num_tokens, num_layers); + dim3 block(std::min(kv_lora_rank, 512)); + vllm::concat_and_cache_mla_grouped_kernel + <<>>( + reinterpret_cast(kv_c.data_ptr()), + reinterpret_cast(k_pe.data_ptr()), + kv_cache_ptrs.const_data_ptr(), + slot_mapping.const_data_ptr(), kv_c_layer_stride, + kv_c_token_stride, k_pe_layer_stride, k_pe_token_stride, + slot_layer_stride, block_stride, entry_stride, kv_lora_rank, pe_dim, + block_size); +} + namespace vllm { template diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu new file mode 100644 index 00000000000..88aed49d181 --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu @@ -0,0 +1,362 @@ +#include "torch_utils.h" + +#include +#include +#include +#include + +#include "custom_all_reduce.cuh" +#include "custom_all_gather_reduce_scatter.cuh" + +namespace vllm { + +void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output, + int size_bytes, int threads, int block_limit) { + if (size_bytes % sizeof(CopyPack) != 0) + throw std::runtime_error( + "custom allgather requires input byte size to be a multiple of " + + std::to_string(sizeof(CopyPack))); + + auto ptrs = buffers_.at(input); + int size_per_rank = size_bytes / sizeof(CopyPack); + int total_size = size_per_rank * world_size_; + int blocks = std::min(block_limit, (total_size + threads - 1) / threads); + +#define AG_CASE(ngpus) \ + case ngpus: \ + cross_device_all_gather<<>>( \ + ptrs, sg_, self_sg_, reinterpret_cast(output), rank_, \ + size_per_rank); \ + break; + + switch (world_size_) { + AG_CASE(2) + AG_CASE(4) + AG_CASE(6) + AG_CASE(8) + default: + throw std::runtime_error( + "custom allgather only supports num gpus in (2,4,6,8)"); + } +#undef AG_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input, + T* output, void* local_buffer, + void* multicast_buffer, + uint32_t* epochs, int size_bytes, + int stage_size_bytes) { + if (size_bytes % sizeof(typename packed_t::P) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport allgather requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size_bytes / sizeof(typename packed_t::P); + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks = + (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportAgThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather, \ + ptrs, input, output, \ + reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ + mnnvl_lamport_all_gather \ + <<>>( \ + ptrs, input, output, reinterpret_cast(multicast_buffer), \ + epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_AG_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_AG_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_AG_CASE(2) + MNNVL_LAMPORT_AG_CASE(4) + MNNVL_LAMPORT_AG_CASE(6) + MNNVL_LAMPORT_AG_CASE(8) + MNNVL_LAMPORT_AG_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_AG_CASE +#undef MNNVL_LAMPORT_AG_LAUNCH +} + +template +void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output, + int size, int threads, int block_limit) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0) + throw std::runtime_error( + "custom reduce-scatter requires each output shard byte size to be " + "a multiple of 16"); + + auto ptrs = buffers_.at(input); + int size_per_rank = size / packed_size / world_size_; + int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads); + +#define RS_CASE(ngpus) \ + case ngpus: \ + cross_device_reduce_scatter<<>>( \ + ptrs, sg_, self_sg_, output, rank_, size_per_rank); \ + break; + + switch (world_size_) { + RS_CASE(2) + RS_CASE(4) + RS_CASE(6) + RS_CASE(8) + default: + throw std::runtime_error( + "custom reduce-scatter only supports num gpus in (2,4,6,8)"); + } +#undef RS_CASE +} + +template +void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream, + T* input, T* output, + void* local_buffer, + uint32_t* epochs, int size, + int stage_size_bytes) { + auto packed_size = packed_t::P::size; + if (size % (packed_size * world_size_) != 0 || + stage_size_bytes % sizeof(typename packed_t::P) != 0) + throw std::runtime_error( + "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes"); + + auto ptrs = buffers_.at(local_buffer); + int size_per_rank = size / packed_size / world_size_; + int stage_size = stage_size_bytes / sizeof(typename packed_t::P); + int blocks_per_rank = + (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads; + int blocks = blocks_per_rank * world_size_; + +#if !defined(USE_ROCM) && CUDA_VERSION >= 12000 + cudaLaunchAttribute attributes[1]{}; + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attributes[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config{.gridDim = dim3(blocks), + .blockDim = dim3(kMnnvlLamportRsThreads), + .dynamicSmemBytes = 0, + .stream = stream, + .attrs = attributes, + .numAttrs = 1}; + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + CUDACHECK(cudaLaunchKernelEx( \ + &config, &mnnvl_lamport_reduce_scatter_kernel, ptrs, input, \ + output, epochs, rank_, size_per_rank, stage_size)) +#else + #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ + mnnvl_lamport_reduce_scatter_kernel \ + <<>>( \ + ptrs, input, output, epochs, rank_, size_per_rank, stage_size) +#endif + +#define MNNVL_LAMPORT_RS_CASE(ngpus) \ + case ngpus: \ + MNNVL_LAMPORT_RS_LAUNCH(ngpus); \ + break; + + switch (world_size_) { + MNNVL_LAMPORT_RS_CASE(2) + MNNVL_LAMPORT_RS_CASE(4) + MNNVL_LAMPORT_RS_CASE(6) + MNNVL_LAMPORT_RS_CASE(8) + MNNVL_LAMPORT_RS_CASE(16) + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports num gpus in " + "(2,4,6,8,16)"); + } +#undef MNNVL_LAMPORT_RS_CASE +#undef MNNVL_LAMPORT_RS_LAUNCH +} + +} // namespace vllm + +using fptr_t = int64_t; +static_assert(sizeof(void*) == sizeof(fptr_t)); + +bool _is_weak_contiguous(torch::stable::Tensor& t); + +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size); +} + +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _local_buffer, + fptr_t _multicast_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto multicast_buffer = reinterpret_cast(_multicast_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_allgather( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + multicast_buffer, epochs, input_size, stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport allgather only supports float32, float16 and " + "bfloat16"); + } +} + +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + auto reg_buffer = reinterpret_cast(_reg_buffer); + STD_TORCH_CHECK(reg_buffer != nullptr); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->reduce_scatter(stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), + inp.numel()); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->reduce_scatter( + stream, reinterpret_cast(reg_buffer), + reinterpret_cast(out.mutable_data_ptr()), inp.numel()); + break; + } +#endif + default: + throw std::runtime_error( + "custom reduce-scatter only supports float32, float16 and bfloat16"); + } +} + +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t _local_buffer, fptr_t _epoch_buffer, + int64_t stage_sz_bytes) { + auto fa = reinterpret_cast(_fa); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); + + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); + auto input_size = inp.numel() * inp.element_size(); + STD_TORCH_CHECK(input_size <= stage_sz_bytes); + auto local_buffer = reinterpret_cast(_local_buffer); + auto epochs = reinterpret_cast(_epoch_buffer); + switch (out.scalar_type()) { + case torch::headeronly::ScalarType::Float: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, epochs, + inp.numel(), stage_sz_bytes); + break; + } +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case torch::headeronly::ScalarType::BFloat16: { + fa->mnnvl_lamport_reduce_scatter( + stream, reinterpret_cast(inp.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), local_buffer, + epochs, inp.numel(), stage_sz_bytes); + break; + } +#endif + default: + throw std::runtime_error( + "MNNVL Lamport reduce-scatter only supports float32, float16 and " + "bfloat16"); + } +} diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp new file mode 100644 index 00000000000..198b78d56da --- /dev/null +++ b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp @@ -0,0 +1,29 @@ +#include "ops.h" +#include "core/registration.h" + +#include + +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) { + custom_ag_rs.def( + "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int multicast_buffer, int epoch_buffer, int " + "stage_sz_bytes) -> ()"); + custom_ag_rs.def( + "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ag_rs.def( + "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int " + "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) { + custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather)); + custom_ag_rs.impl("mnnvl_lamport_all_gather", + TORCH_BOX(&mnnvl_lamport_all_gather)); + custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter)); + custom_ag_rs.impl("mnnvl_lamport_reduce_scatter", + TORCH_BOX(&mnnvl_lamport_reduce_scatter)); +} diff --git a/csrc/libtorch_stable/custom_all_reduce.cu b/csrc/libtorch_stable/custom_all_reduce.cu index 0f7f759949a..07fca99ddd9 100644 --- a/csrc/libtorch_stable/custom_all_reduce.cu +++ b/csrc/libtorch_stable/custom_all_reduce.cu @@ -18,14 +18,14 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, torch::stable::Tensor& rank_data, int64_t rank, bool fully_connected) { int world_size = fake_ipc_ptrs.size(); - if (world_size > 8) - throw std::invalid_argument("world size > 8 is not supported"); + if (world_size > vllm::kMaxCustomCollectiveRanks) + throw std::invalid_argument("world size > 16 is not supported"); if (world_size % 2 != 0) throw std::invalid_argument("Odd num gpus is not supported for now"); if (rank < 0 || rank >= world_size) throw std::invalid_argument("invalid rank passed in"); - vllm::Signal* ipc_ptrs[8]; + vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks]; for (int i = 0; i < world_size; i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } @@ -124,7 +124,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); } void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); - void* ipc_ptrs[8]; + void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks]; for (int i = 0; i < fake_ipc_ptrs.size(); i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } diff --git a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu index 585004c047b..d87b034be8c 100644 --- a/csrc/libtorch_stable/dsv3_fused_a_gemm.cu +++ b/csrc/libtorch_stable/dsv3_fused_a_gemm.cu @@ -647,17 +647,17 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( #endif } -template +template void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, - cudaStream_t const stream) { - constexpr int gemm_m = kHdOut; // 2112 - int const gemm_n = num_tokens; // 1-16 - constexpr int gemm_k = kHdIn; // 7168 + cudaStream_t const stream, bool enable_pdl) { + constexpr int gemm_m = kHdOut; + int const gemm_n = num_tokens; + constexpr int gemm_k = kHdIn; constexpr int batch_size = 1; std::swap(mat_a, mat_b); constexpr int tile_m = 16; - constexpr int tile_n = kTileN; // 8 or 16 - constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 + constexpr int tile_n = kTileN; + constexpr int tile_k = kTileK; constexpr int max_stage_cnt = 1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t)); constexpr int k_iter_cnt = gemm_k / tile_k; @@ -679,7 +679,8 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, config.stream = stream; cudaLaunchAttribute attrs[1]; attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); + attrs[0].val.programmaticStreamSerializationAllowed = + enable_pdl || getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; if (smem_bytes >= (48 * 1024)) { @@ -694,36 +695,50 @@ void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, output, mat_a, mat_b, gemm_n); } -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, - cudaStream_t); +template +void invokeFusedAGemmForTokens(T* output, T const* mat_a, T const* mat_b, + int num_tokens, cudaStream_t const stream, + bool enable_pdl) { + if (num_tokens <= 8) { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } else { + invokeFusedAGemm( + output, mat_a, mat_b, num_tokens, stream, enable_pdl); + } +} void dsv3_fused_a_gemm(torch::stable::Tensor& output, torch::stable::Tensor const& mat_a, - torch::stable::Tensor const& mat_b) { + torch::stable::Tensor const& mat_b, bool enable_pdl) { STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2); int const num_tokens = mat_a.size(0); int const hd_in = mat_a.size(1); int const hd_out = mat_b.size(1); - constexpr int kHdIn = 7168; - constexpr int kHdOut = 2112; STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, "required 1 <= mat_a.shape[0] <= 16"); - STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168"); - STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112"); STD_TORCH_CHECK(output.size(0) == num_tokens, "required output.shape[0] == mat_a.shape[0]"); STD_TORCH_CHECK(output.size(1) == hd_out, "required output.shape[1] == mat_b.shape[1]"); + STD_TORCH_CHECK(mat_b.size(0) == hd_in, + "required mat_b.shape[0] == mat_a.shape[1]"); - STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); - STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor"); - STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); + STD_TORCH_CHECK(mat_a.get_device_index() == mat_b.get_device_index() && + mat_a.get_device_index() == output.get_device_index(), + "mat_a, mat_b, and output must be on the same device"); + + // The kernels index global memory with raw pointers and packed strides, so + // reject any padded or transposed view rather than reading out of bounds. + STD_TORCH_CHECK( + mat_a.stride(0) == hd_in && mat_a.stride(1) == 1, + "mat_a must be a packed row-major [num_tokens, hd_in] tensor"); + STD_TORCH_CHECK( + output.stride(0) == hd_out && output.stride(1) == 1, + "output must be a packed row-major [num_tokens, hd_out] tensor"); + STD_TORCH_CHECK(mat_b.stride(0) == 1 && mat_b.stride(1) == hd_in, + "mat_b must be a packed column-major [hd_in, hd_out] tensor"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 && @@ -738,19 +753,85 @@ void dsv3_fused_a_gemm(torch::stable::Tensor& output, STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90"); auto stream = get_current_cuda_stream(mat_a.get_device_index()); - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens, - stream); + auto* output_ptr = + reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()); + auto const* mat_a_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); + auto const* mat_b_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()); + +#define DISPATCH_DSV3_SHAPE(HD_IN, HD_OUT) \ + if (hd_in == HD_IN && hd_out == HD_OUT) { \ + invokeFusedAGemmForTokens<__nv_bfloat16, HD_IN, HD_OUT>( \ + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); \ + return; \ } + + // Shapes the Kimi-K3 selector routes to dsv3_fused_a (see the dsv3 winners + // in KIMI_K3_PROJECTIONS) plus the DeepSeek V2/V3 QKV A-projection. + DISPATCH_DSV3_SHAPE(7168, 1536) + DISPATCH_DSV3_SHAPE(7168, 2112) + DISPATCH_DSV3_SHAPE(1536, 2304) + DISPATCH_DSV3_SHAPE(1536, 4608) + DISPATCH_DSV3_SHAPE(7168, 3584) + DISPATCH_DSV3_SHAPE(768, 7168) + // TP16 dsv3 winners, as (hd_in=K, hd_out=N). TP16 dense down_proj is absent + // because hd_in=2112 is not a multiple of any supported tile_k. + DISPATCH_DSV3_SHAPE(1536, 1152) + DISPATCH_DSV3_SHAPE(7168, 768) + DISPATCH_DSV3_SHAPE(7168, 3216) + DISPATCH_DSV3_SHAPE(7168, 4224) + +#ifdef VLLM_K3_BENCH_SHAPES + // The selector routes these shapes to CuTe or the default GEMM, so they are + // never reached in production. They are compiled only for offline + // DSV3-vs-CuTe benchmarking. + DISPATCH_DSV3_SHAPE(7168, 6288) + DISPATCH_DSV3_SHAPE(1536, 7168) + DISPATCH_DSV3_SHAPE(3584, 7168) + DISPATCH_DSV3_SHAPE(7168, 8448) + DISPATCH_DSV3_SHAPE(7168, 20480) + DISPATCH_DSV3_SHAPE(7168, 3072) + DISPATCH_DSV3_SHAPE(7168, 12448) + DISPATCH_DSV3_SHAPE(3072, 7168) + DISPATCH_DSV3_SHAPE(8448, 7168) + DISPATCH_DSV3_SHAPE(7168, 16896) + DISPATCH_DSV3_SHAPE(7168, 40960) +#endif + +#undef DISPATCH_DSV3_SHAPE + + if (hd_in == 128 && hd_out == 1536) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 1536, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 128 && hd_out == 3072) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 3072, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + // TP16 KDA f_b_proj and shared_expert down_proj. Neither hd_in is a multiple + // of 256, so both need the 128 tile_k. + if (hd_in == 128 && hd_out == 768) { + invokeFusedAGemmForTokens<__nv_bfloat16, 128, 768, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } + if (hd_in == 384 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 384, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#ifdef VLLM_K3_BENCH_SHAPES + if (hd_in == 4224 && hd_out == 7168) { + invokeFusedAGemmForTokens<__nv_bfloat16, 4224, 7168, 128>( + output_ptr, mat_a_ptr, mat_b_ptr, num_tokens, stream, enable_pdl); + return; + } +#endif + + STD_TORCH_CHECK(false, "unsupported DSV3 fused-A GEMM shape"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu new file mode 100644 index 00000000000..cbec4ff71a6 --- /dev/null +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -0,0 +1,1237 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + * + * Fused Kimi-K3 MLA prefill + decode epilogues with optional RoPE. + * + * Prefill: runs after the q_b_proj / kv_b_proj GEMMs, one launch per token + * slice. Decode: runs after BMM1 (q_nope x W_UK) right before forward_mqa, + * concatenating mqa_q = [ql_nope | q_pe] and inserting the latent cache + * (fused_kimi_k3_mla_decode_q_concat_kv_cache_{,fp8_,ds_mla_}insert). + * + * Prefill variants: + * + * bf16 (fused_kimi_k3_mla_key_concat_kv_cache_insert): + * - optional in-place q RoPE: rotate q[t, h, 128:192] + * - full key concat: k_out[t, h] = [k_nope[t, h] | k_pe[t]] (per head) + * - latent cache insert: cache[slot(t)] = [kv_c_normed[t] | k_pe[t]] + * (v is used as-is in bf16, so it is not touched here.) + * + * fp8 (fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert): + * - q_fp8[t, h] = quant(q[t, h], q_scale) + * - k_fp8[t, h] = quant([k_nope[t, h] | k_pe[t]], k_scale) + * - v_fp8[t, h] = quant(v[t, h], v_scale) + * - cache[slot(t)] = quant([kv_c_normed[t] | k_pe[t]], k_scale) + * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a + * separate cache scale). Per-tensor E4M3. + * + * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): + * - full key concat (bf16), and cache insert in DeepSeek's 656-byte + * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic + * scales, RoPE bf16), bit-compatible with concat_and_cache_ds_mla_kernel. + * + * Both use Programmatic Dependent Launch (PDL) to overlap the tail of the + * producing GEMMs on sm_90+, and are structured after + * `fusedDeepseekV4FullCacheKernel`: one grid, one warp per (token, slot) with + * `slotsPerToken = num_heads + 1`. Slots [0, H) do the per-head work; the extra + * slot H does the per-token cache insert. All dims are multiples of 8, so bf16 + * copies move one uint4 (8 elems) per step and fp8 stores pack 8 elems into a + * uint2. + */ + +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + +#include "cuda_compat.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" + +#ifndef USE_ROCM + #include + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" +#else + #include + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#endif +#include +#include +#include + +#ifdef USE_ROCM +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { + #if defined(__gfx942__) + __hip_fp8_e4m3_fnuz fp8_val(val); + #else + __hip_fp8_e4m3 fp8_val(val); + #endif + return reinterpret_cast(fp8_val); +} +#endif + +namespace vllm { +namespace kimi_k3_fused_ops { + +namespace { +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} +} // namespace + +// ──────────────────────────────────────────────────────────────────────────── +// Constants (Kimi-K3 MLA) +// ──────────────────────────────────────────────────────────────────────────── +constexpr int kKvLoraRank = 512; // L +constexpr int kQkNopeHeadDim = 128; // P +constexpr int kQkRopeHeadDim = 64; // R +constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192 +constexpr int kVHeadDim = 128; // V +constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576 +constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store + +#if defined(USE_ROCM) && defined(__gfx942__) +constexpr float kFp8Max = 224.0f; +#else +constexpr float kFp8Max = 448.0f; +#endif +// Divisor for fp8_ds_mla per-tile dynamic scales (matches cache_kernels.cu). +// fp8_ds_mla 656B entry: [0,512) NoPE fp8 (4 tiles of 128), [512,528) 4 fp32 +// tile scales, [528,656) RoPE 64 bf16. +constexpr float kFp8ScaleDivisor = kFp8Max; + +// Copy 8 source elements (one uint4 of bf16/fp16) to `dst`. FP8=false stores a +// uint4 (bf16); FP8=true decodes to fp32, scales by `scale_inv`, saturates to +// ±kFp8Max and packs into a uint2 of E4M3. +template +__device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, + float scale_inv, + const float* cos_sin = nullptr, + int rope_elem_base = 0) { + uint4 const v = *reinterpret_cast(src); + if constexpr (FP8 || APPLY_ROPE) { + using Converter = vllm::_typeConvert; + auto const* p = + reinterpret_cast(&v); + float f[kVecElems]; +#pragma unroll + for (int i = 0; i < 4; i++) { + float2 x = Converter::convert(p[i]); + f[2 * i] = x.x; + f[2 * i + 1] = x.y; + } + if constexpr (APPLY_ROPE) { +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + int const pair_idx = rope_elem_base / 2 + i; + float const cos = static_cast(cos_sin[pair_idx]); + float const sin = + static_cast(cos_sin[pair_idx + kQkRopeHeadDim / 2]); + float const x = f[2 * i]; + float const y = f[2 * i + 1]; + f[2 * i] = x * cos - y * sin; + f[2 * i + 1] = x * sin + y * cos; + } + } + if constexpr (!FP8) { + uint4 out; + auto* o = reinterpret_cast(&out); +#pragma unroll + for (int i = 0; i < kVecElems / 2; i++) { + o[i] = Converter::convert(make_float2(f[2 * i], f[2 * i + 1])); + } + *reinterpret_cast(dst) = out; + return; + } +#ifndef USE_ROCM + uint2 out; + auto* o2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out); + #pragma unroll + for (int i = 0; i < 4; i++) { + float2 s = make_float2(f[2 * i] * scale_inv, f[2 * i + 1] * scale_inv); + s.x = fminf(fmaxf(s.x, -kFp8Max), kFp8Max); + s.y = fminf(fmaxf(s.y, -kFp8Max), kFp8Max); + o2[i] = __nv_cvt_float2_to_fp8x2(s, __NV_SATFINITE, __NV_E4M3); + } + *reinterpret_cast(dst) = out; +#else + uint8_t out[kVecElems]; + #pragma unroll + for (int i = 0; i < kVecElems; i++) { + float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max); + out[i] = rocm_cvt_float_to_fp8_e4m3(s); + } + *reinterpret_cast(dst) = *reinterpret_cast(out); +#endif + } else { + *reinterpret_cast(dst) = v; + } +} + +// Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). +// FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size). +template +__device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope, + const scalar_t* k_pe, int laneId, + int dst_elem_size, float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, k_nope + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8( + d + e * dst_elem_size, k_pe + rope_e, scale_inv, cos_sin, rope_e); + } + } +} + +// Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias +// q (in-place); fp8 writes the quantized query directly to its output. +template +__device__ __forceinline__ void writePrefillQuery( + void* dst, const scalar_t* q, int laneId, int dst_elem_size, + float scale_inv, const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + if constexpr (FP8) { + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 32 * kVecElems) { + if (e < kQkNopeHeadDim) { + copyChunk8(d + e * dst_elem_size, q + e, scale_inv); + } else { + int const rope_e = e - kQkNopeHeadDim; + copyChunk8(d + e * dst_elem_size, q + e, + scale_inv, cos_sin, rope_e); + } + } + } else if constexpr (APPLY_ROPE) { + for (int e = laneId * kVecElems; e < kQkRopeHeadDim; e += 32 * kVecElems) { + copyChunk8( + d + (kQkNopeHeadDim + e) * dst_elem_size, q + kQkNopeHeadDim + e, + scale_inv, cos_sin, e); + } + } +} + +// Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used +// for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache +// entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4. +template +__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512, + const scalar_t* b64, int laneId, + int dst_elem_size, + float scale_inv, + const float* cos_sin = nullptr) { + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) { + if (e < kKvLoraRank) { + copyChunk8(d + e * dst_elem_size, a512 + e, scale_inv); + } else { + int const rope_e = e - kKvLoraRank; + copyChunk8(d + e * dst_elem_size, b64 + rope_e, + scale_inv, cos_sin, rope_e); + } + } +} + +// Write [kv_c | k_pe] into the fp8_ds_mla 656B entry using one warp: NoPE 512 +// as fp8 in 4 tiles of 128 (per-tile dynamic absmax scale, 4 fp32 scales at +// [512,528)), RoPE 64 as bf16 at [528,656). Bit-compatible with +// concat_and_cache_ds_mla_kernel. +template +__device__ __forceinline__ void writeDsMlaCache( + uint8_t* row, const scalar_t* kvc, const scalar_t* pe, int laneId, + const float* cos_sin = nullptr) { + constexpr int kElemsPerLane = kKvLoraRank / 32; // 16 + int const tile = laneId >> 3; // 8 lanes per tile + scalar_t vals[kElemsPerLane]; + *reinterpret_cast(vals) = + *reinterpret_cast(kvc + laneId * kElemsPerLane); + *reinterpret_cast(vals + 8) = + *reinterpret_cast(kvc + laneId * kElemsPerLane + 8); + + float max_abs = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + max_abs = fmaxf(max_abs, fabsf(static_cast(vals[i]))); + } +#pragma unroll + for (int offset = 4; offset > 0; offset /= 2) { + max_abs = fmaxf(max_abs, VLLM_SHFL_XOR_SYNC_WIDTH(max_abs, offset, 8)); + } + float const tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN); + if ((laneId & 7) == 0) { + reinterpret_cast(row)[kKvLoraRank / 4 + tile] = tile_scale; + } + uint8_t res[kElemsPerLane]; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + res[i] = + fp8::scaled_convert( + vals[i], tile_scale); + } + *reinterpret_cast(row + laneId * kElemsPerLane) = + *reinterpret_cast(res); + scalar_t* row16 = reinterpret_cast(row); + scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2; + if constexpr (APPLY_ROPE) { + using Converter = vllm::_typeConvert; + using packed_t = typename Converter::packed_hip_type; + packed_t const src = *reinterpret_cast(pe + laneId * 2); + float2 const xy = Converter::convert(src); + float const cos = static_cast(cos_sin[laneId]); + float const sin = static_cast(cos_sin[laneId + kQkRopeHeadDim / 2]); + *reinterpret_cast(rope_dst) = Converter::convert( + make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos)); + } else { + *reinterpret_cast(rope_dst) = + *reinterpret_cast(pe + laneId * 2); + } +} + +// ──────────────────────────────────────────────────────────────────────────── +// bf16 variant: optional q RoPE + full-key concat + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, scalar_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache, int const num_tokens, + int const num_heads, int const cache_block_size) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + scalar_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// fp8 variant: quant q / k / v + latent cache insert +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( + const scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ v, int64_t const v_tok_stride, + int64_t const v_head_stride, uint8_t* __restrict__ q_fp8, + int64_t const qo_tok_stride, int64_t const qo_head_stride, + uint8_t* __restrict__ k_fp8, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8, + int64_t const vo_tok_stride, int64_t const vo_head_stride, + uint8_t* __restrict__ k_cache, int64_t const cache_block_stride, + int64_t const cache_token_stride, const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ k_scale_inv, + const float* __restrict__ v_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + int const h = slotIdx; + // q_fp8[t, h] = quant(q[t, h], q_scale) + float const qsi = __ldg(q_scale_inv); + const scalar_t* qh = q + tokenIdx * q_tok_stride + h * q_head_stride; + uint8_t* qo = q_fp8 + tokenIdx * qo_tok_stride + h * qo_head_stride; + writePrefillQuery(qo, qh, laneId, 1, qsi, + rope_cache); + // k_fp8[t, h] = quant([k_nope | k_pe], k_scale) + writeFullKey( + k_fp8 + tokenIdx * ko_tok_stride + h * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + h * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, __ldg(k_scale_inv), + rope_cache); + // v_fp8[t, h] = quant(v[t, h], v_scale) + float const vsi = __ldg(v_scale_inv); + const scalar_t* vh = v + tokenIdx * v_tok_stride + h * v_head_stride; + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + h * vo_head_stride; + for (int e = laneId * kVecElems; e < kVHeadDim; e += 32 * kVecElems) { + copyChunk8(vo + e, vh + e, vsi); + } + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + // The cache latent uses _k_scale (read back by decode / context); the + // attention key (k_fp8 above) uses its own k_scale. + float const ksi = __ldg(cache_scale_inv); + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeLatent576( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert +// +// Cache entry (656 bytes), matching concat_and_cache_ds_mla_kernel: +// [0, 512) NoPE 512 vals as fp8, 4 tiles of 128, each dynamically scaled +// [512, 528) 4 fp32 per-tile scales +// [528, 656) RoPE 64 vals as bf16 (unquantized) +// The cache slot uses one warp: lane L quantizes NoPE elems [L*16, L*16+16) +// (tile = L>>3, absmax-reduced within its 8-lane tile group), then all 32 lanes +// write 2 RoPE bf16 each. +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel( + scalar_t* __restrict__ q, int64_t const q_tok_stride, + int64_t const q_head_stride, const scalar_t* __restrict__ k_nope, + int64_t const kn_tok_stride, int64_t const kn_head_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + scalar_t* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + scalar_t* qh = q + tokenIdx * q_tok_stride + slotIdx * q_head_stride; + writePrefillQuery( + qh, qh, laneId, sizeof(scalar_t), 1.0f, rope_cache); + // Full key (bf16): k_out[t, h] = [k_nope[t, h] | k_pe[t]]. + writeFullKey( + k_out + tokenIdx * ko_tok_stride + slotIdx * ko_head_stride, + k_nope + tokenIdx * kn_tok_stride + slotIdx * kn_head_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert, +// run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the +// plain per-tensor cache. (ds_mla cache uses the separate kernel below.) +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + void* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, void* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, + const float* __restrict__ q_scale_inv, + const float* __restrict__ cache_scale_inv, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + constexpr int kMqElem = Q_FP8 ? 1 : sizeof(scalar_t); + constexpr int kCacheElem = KV_FP8 ? 1 : sizeof(scalar_t); + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(mqa_q) + + (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + kMqElem, qsi, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + float const ksi = KV_FP8 ? __ldg(cache_scale_inv) : 1.0f; + writeLatent576( + reinterpret_cast(k_cache) + + (slot_id / cache_block_size * cache_block_stride + + slot_id % cache_block_size * cache_token_stride) * + kCacheElem, + kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride, + laneId, kCacheElem, ksi, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// Decode epilogue for fp8_ds_mla: concat mqa_q (bf16) + ds_mla cache insert. +template +__global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( + const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, + int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, + int64_t const qpe_tok_stride, int64_t const qpe_head_stride, + const scalar_t* __restrict__ kv_c, int64_t const kv_c_tok_stride, + const scalar_t* __restrict__ k_pe, int64_t const k_pe_tok_stride, + scalar_t* __restrict__ mqa_q, int64_t const mq_tok_stride, + int64_t const mq_head_stride, uint8_t* __restrict__ k_cache, + int64_t const cache_block_stride, int64_t const cache_token_stride, + const int64_t* __restrict__ slot_mapping, int const num_tokens, + int const num_heads, int const cache_block_size, + const int64_t* __restrict__ position_ids, + const float* __restrict__ cos_sin_cache) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; + int const slotsPerToken = num_heads + 1; + int const tokenIdx = globalWarpIdx / slotsPerToken; + int const slotIdx = globalWarpIdx % slotsPerToken; + if (tokenIdx >= num_tokens) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + const float* rope_cache = nullptr; + if constexpr (APPLY_ROPE) { + rope_cache = cos_sin_cache + position_ids[tokenIdx] * kQkRopeHeadDim; + } + + if (slotIdx < num_heads) { + writeLatent576( + mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride, + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, + sizeof(scalar_t), 1.0f, rope_cache); + } else { + int64_t const slot_id = slot_mapping[tokenIdx]; + if (slot_id >= 0) { + uint8_t* row = k_cache + + (slot_id / cache_block_size) * cache_block_stride + + (slot_id % cache_block_size) * cache_token_stride; + writeDsMlaCache( + row, kv_c + tokenIdx * kv_c_tok_stride, + k_pe + tokenIdx * k_pe_tok_stride, laneId, rope_cache); + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +// PDL-aware launch of a (token, num_heads + 1)-warp grid. +template +static void launchPdl(KernelT kernel, int num_tokens, int num_heads, + cudaStream_t stream, Args... args) { + constexpr int kBlockSize = 256; + constexpr int kWarpsPerBlock = kBlockSize / 32; + int64_t const total_warps = + static_cast(num_tokens) * (num_heads + 1); + int const grid = + static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); +#ifndef USE_ROCM + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3(grid); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + cudaLaunchKernelEx(&config, kernel, args...); +#else + // clang-format off + // hipify's CUDA->HIP regex catastrophically backtracks on "> > >"; keep the + // launch closer as ">>>". clang-format would otherwise re-split it (it does + // not parse the CUDA launch syntax in this libtorch_stable file). + kernel<<>>(args...); + // clang-format on +#endif +} + +} // namespace kimi_k3_fused_ops +} // namespace vllm + +// ──────────────────────────────────────────────────────────────────────────── +// Torch op wrappers +// ──────────────────────────────────────────────────────────────────────────── +namespace { +bool check_rope_inputs( + std::optional const& position_ids, + std::optional const& cos_sin_cache, + torch::stable::Tensor const& /*input*/, int64_t num_tokens) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(position_ids.has_value() == cos_sin_cache.has_value(), + "position_ids and cos_sin_cache must be provided together"); + if (!position_ids.has_value()) return false; + + auto const& positions = position_ids.value(); + auto const& rope_cache = cos_sin_cache.value(); + STD_TORCH_CHECK(positions.device().is_cuda() && positions.dim() == 1 && + positions.scalar_type() == ScalarType::Long && + positions.size(0) == num_tokens, + "position_ids must be int64 CUDA with shape [num_tokens]"); + STD_TORCH_CHECK(rope_cache.device().is_cuda() && rope_cache.dim() == 2 && + rope_cache.size(1) == 64 && rope_cache.stride(1) == 1 && + rope_cache.scalar_type() == ScalarType::Float, + "cos_sin_cache must have shape [max_position, 64], unit " + "last-dim stride, and be fp32 (RoPE math runs in fp32)"); + return true; +} +} // namespace + +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] + torch::stable::Tensor const& k_pe, // [Tp, 64] + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192] CUDA"); + // k_pe is a strided view of the fused QKV-LoRA GEMM output; the kernel takes + // its row stride and reads the 64 cols contiguously, so unit last-dim stride + // (not full contiguity) is all that is required. + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride, CUDA"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && kv_c_normed.dim() == 2 && + kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous CUDA"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.dim() == 3 && k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous CUDA"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous CUDA"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK(q.scalar_type() == dt && k_pe.scalar_type() == dt && + kv_c_normed.scalar_type() == dt && + k_out.scalar_type() == dt && k_cache.scalar_type() == dt, + "all tensors must share k_nope's (bf16/fp16) dtype"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr, + num_tokens, num_heads, static_cast(cache_block_size)); + }; + if (apply_rope) { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatKVCacheInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, // [Tp, H, 192] + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor& k_out, // [Tp, H, 192] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK( + k_nope.device().is_cuda() && k_nope.dim() == 3 && k_nope.size(2) == 128, + "k_nope shape [Tp, H, 128] CUDA"); + STD_TORCH_CHECK(q.device().is_cuda() && q.scalar_type() == dt && + q.dim() == 3 && q.size(2) == 192, + "q shape [Tp, H, 192]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + STD_TORCH_CHECK(k_out.device().is_cuda() && k_out.is_contiguous() && + k_out.scalar_type() == dt && k_out.dim() == 3 && + k_out.size(2) == 192, + "k_out shape [Tp, H, 192] contiguous"); + // fp8_ds_mla entry is 656 bytes stored as uint8. + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + STD_TORCH_CHECK(static_cast(k_out.size(1)) == num_heads, + "k_out head count must match k_nope"); + STD_TORCH_CHECK(q.size(0) == num_tokens && q.size(1) == num_heads, + "q token/head dimensions must match k_nope"); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_key_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.mutable_data_ptr()), q.stride(0), + q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_out.mutable_data_ptr()), + k_out.stride(0), k_out.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } else { + launch( + kk3::fusedKimiK3MLAKeyConcatDsMlaInsertKernel); + } + }); +} + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, // [Tp, H, 192] bf16 + torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 + torch::stable::Tensor const& k_pe, // [Tp, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [Tp, 512] bf16 + torch::stable::Tensor const& v, // [Tp, H, 128] bf16 + torch::stable::Tensor& q_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& k_fp8, // [Tp, H, 192] fp8, written + torch::stable::Tensor& v_fp8, // [Tp, H, 128] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [Tp] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& k_scale_inv, // scalar fp32 (1 / k scale) + torch::stable::Tensor const& v_scale_inv, // scalar fp32 (1 / v scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + auto check_in = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.scalar_type() == dt && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_in(q, 192, "q shape [Tp, H, 192]"); + check_in(k_nope, 128, "k_nope shape [Tp, H, 128]"); + check_in(v, 128, "v shape [Tp, H, 128]"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [Tp, 64], unit last-dim stride"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [Tp, 512] contiguous"); + auto check_out = [&](torch::stable::Tensor const& t, int d2, char const* n) { + STD_TORCH_CHECK(t.device().is_cuda() && t.is_contiguous() && + t.scalar_type() == ScalarType::Float8_e4m3fn && + t.dim() == 3 && t.size(2) == d2, + n); + }; + check_out(q_fp8, 192, "q_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(k_fp8, 192, "k_fp8 shape [Tp, H, 192] fp8 contiguous"); + check_out(v_fp8, 128, "v_fp8 shape [Tp, H, 128] fp8 contiguous"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(k_scale_inv, "k_scale_inv must be scalar float32 CUDA"); + check_scale(v_scale_inv, "v_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q, num_tokens); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(q.const_data_ptr()), + q.stride(0), q.stride(1), + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(v.const_data_ptr()), + v.stride(0), v.stride(1), + reinterpret_cast(q_fp8.mutable_data_ptr()), + q_fp8.stride(0), q_fp8.stride(1), + reinterpret_cast(k_fp8.mutable_data_ptr()), + k_fp8.stride(0), k_fp8.stride(1), + reinterpret_cast(v_fp8.mutable_data_ptr()), + v_fp8.stride(0), v_fp8.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + k_scale_inv.const_data_ptr(), + v_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } else { + launch(kk3::fusedKimiK3MLAQKVQuantKVCacheFp8Kernel); + } + }); +} + +// ──────────────────────────────────────────────────────────────────────────── +// Decode epilogue torch ops +// ──────────────────────────────────────────────────────────────────────────── +namespace { +// Shared shape checks for the decode ops. Verifies the query/latent inputs and +// slot_mapping; caller checks mqa_q / k_cache dtypes for its variant. +void check_decode_inputs(torch::stable::Tensor const& ql_nope, + torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& mqa_q, + torch::stable::Tensor const& slot_mapping) { + using torch::headeronly::ScalarType; + auto const dt = ql_nope.scalar_type(); + STD_TORCH_CHECK(ql_nope.device().is_cuda() && ql_nope.dim() == 3 && + ql_nope.size(2) == 512, + "ql_nope shape [B, H, 512] CUDA"); + STD_TORCH_CHECK(q_pe.device().is_cuda() && q_pe.scalar_type() == dt && + q_pe.dim() == 3 && q_pe.size(2) == 64, + "q_pe shape [B, H, 64]"); + STD_TORCH_CHECK(kv_c_normed.device().is_cuda() && + kv_c_normed.is_contiguous() && + kv_c_normed.scalar_type() == dt && + kv_c_normed.dim() == 2 && kv_c_normed.size(1) == 512, + "kv_c_normed shape [B, 512] contiguous"); + STD_TORCH_CHECK(k_pe.device().is_cuda() && k_pe.dim() == 2 && + k_pe.stride(1) == 1 && k_pe.scalar_type() == dt && + k_pe.size(1) == 64, + "k_pe shape [B, 64], unit last-dim stride"); + STD_TORCH_CHECK(mqa_q.device().is_cuda() && mqa_q.is_contiguous() && + mqa_q.dim() == 3 && mqa_q.size(2) == 576, + "mqa_q shape [B, H, 576] contiguous"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == ScalarType::Long, + "slot_mapping must be int64 CUDA"); +} +} // namespace + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] bf16, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt && k_cache.scalar_type() == dt, + "mqa_q / k_cache must match ql_nope dtype (bf16)"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 576] contiguous"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + nullptr, nullptr, num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] fp8, written + torch::stable::Tensor& k_cache, // [nblk, bs, 576] fp8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + torch::stable::Tensor const& q_scale_inv, // scalar fp32 (1 / q scale) + torch::stable::Tensor const& cache_scale_inv, // scalar fp32 (1 / kv scale) + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == ScalarType::Float8_e4m3fn, + "mqa_q must be float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.device().is_cuda() && k_cache.dim() == 3 && + k_cache.size(1) == cache_block_size && + k_cache.size(2) == 576 && k_cache.stride(2) == 1 && + k_cache.scalar_type() == ScalarType::Float8_e4m3fn, + "k_cache shape [nblk, block_size, 576] fp8 contiguous"); + auto check_scale = [&](torch::stable::Tensor const& s, char const* n) { + STD_TORCH_CHECK(s.device().is_cuda() && + s.scalar_type() == ScalarType::Float && s.size(0) == 1, + n); + }; + check_scale(q_scale_inv, "q_scale_inv must be scalar float32 CUDA"); + check_scale(cache_scale_inv, "cache_scale_inv must be scalar float32 CUDA"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), mqa_q.mutable_data_ptr(), mqa_q.stride(0), + mqa_q.stride(1), k_cache.mutable_data_ptr(), k_cache.stride(0), + k_cache.stride(1), slot_mapping.const_data_ptr(), + q_scale_inv.const_data_ptr(), + cache_scale_inv.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + } + }); +} + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, // [B, H, 512] bf16 + torch::stable::Tensor const& q_pe, // [B, H, 64] bf16 + torch::stable::Tensor const& kv_c_normed, // [B, 512] bf16 + torch::stable::Tensor const& k_pe, // [B, 64] bf16 + torch::stable::Tensor& mqa_q, // [B, H, 576] bf16, written + torch::stable::Tensor& k_cache, // [nblk, bs, 656] uint8, written + torch::stable::Tensor const& slot_mapping, // [B] int64 + int64_t cache_block_size, std::optional position_ids, + std::optional cos_sin_cache) { + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = ql_nope.scalar_type(); + check_decode_inputs(ql_nope, q_pe, kv_c_normed, k_pe, mqa_q, slot_mapping); + STD_TORCH_CHECK(mqa_q.scalar_type() == dt, "mqa_q must be bf16 for ds_mla"); + STD_TORCH_CHECK( + k_cache.device().is_cuda() && k_cache.scalar_type() == ScalarType::Byte && + k_cache.dim() == 3 && k_cache.size(1) == cache_block_size && + k_cache.size(2) == 656 && k_cache.stride(2) == 1, + "k_cache shape [nblk, block_size, 656] uint8 contiguous"); + + int const num_tokens = static_cast(ql_nope.size(0)); + int const num_heads = static_cast(ql_nope.size(1)); + bool const apply_rope = + check_rope_inputs(position_ids, cos_sin_cache, q_pe, num_tokens); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(ql_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdl( + kernel, num_tokens, num_heads, stream, + reinterpret_cast(ql_nope.const_data_ptr()), + ql_nope.stride(0), ql_nope.stride(1), + reinterpret_cast(q_pe.const_data_ptr()), + q_pe.stride(0), q_pe.stride(1), + reinterpret_cast(kv_c_normed.const_data_ptr()), + kv_c_normed.stride(0), + reinterpret_cast(k_pe.const_data_ptr()), + k_pe.stride(0), + reinterpret_cast(mqa_q.mutable_data_ptr()), + mqa_q.stride(0), mqa_q.stride(1), + reinterpret_cast(k_cache.mutable_data_ptr()), + k_cache.stride(0), k_cache.stride(1), + slot_mapping.const_data_ptr(), num_tokens, num_heads, + static_cast(cache_block_size), + apply_rope ? position_ids.value().const_data_ptr() + : nullptr, + apply_rope ? reinterpret_cast( + cos_sin_cache.value().const_data_ptr()) + : nullptr); + }; + if (apply_rope) { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatDsMlaKernel); + } + }); +} diff --git a/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu new file mode 100644 index 00000000000..0badeb9b772 --- /dev/null +++ b/csrc/libtorch_stable/kimi_k3/fused_kda_decode_kernel.cu @@ -0,0 +1,1130 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + */ + +#include +#include +#include +#include +#include + +#include "../torch_utils.h" +#include "../../cuda_compat.h" + +namespace { + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kKernelWidth = 4; +constexpr int kConvStateWidth = kKernelWidth - 1; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; + +struct KdaDecodeStrides { + int64_t x_row; + int64_t beta_row; + int64_t onorm_row; + int64_t conv_slot; + int64_t state_slot; +}; + +__device__ __forceinline__ float bf16_load(const __nv_bfloat16* ptr, + int64_t idx) { + return __bfloat162float(ptr[idx]); +} + +__device__ __forceinline__ float bf16_load(const float* ptr, int64_t idx) { + return ptr[idx]; +} + +template +__device__ __forceinline__ float conv_weight_load(const float* ptr, int channel, + int width) { + return ptr[width * kChannels + channel]; +} + +__device__ __forceinline__ __nv_bfloat16 bf16_store(float value) { + return __float2bfloat16(value); +} + +template +__device__ __forceinline__ void store_state_float4(float* ptr, float4 value) { + if constexpr (kUseCacheGlobalStore) { + __stcg(reinterpret_cast(ptr), value); + } else { + *reinterpret_cast(ptr) = value; + } +} + +__device__ __forceinline__ float sigmoid_fast(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float silu_fast(float x) { + return x * sigmoid_fast(x); +} + +__device__ __forceinline__ float softplus_fast(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} + +__device__ __forceinline__ float warp_reduce_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +__device__ __forceinline__ void cp_async_cg_16b(float* smem_ptr, + const float* gmem_ptr) { + uint32_t smem_addr = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem_addr), "l"(gmem_ptr)); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::); +} + +__device__ __forceinline__ void cp_async_wait_all() { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +} + +__device__ __forceinline__ void cp_async_wait_group_1() { + asm volatile("cp.async.wait_group 1;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cp_async_state_chunk_for(float* s_state, + const float* state, + int slot, int i_hv, + int HV, int chunk) { + constexpr int kFloat4PerChunk = kChunkV * kDimK / 4; + const int tid = threadIdx.x; + const int stage = chunk & 1; + const int v_base = chunk * kChunkV; + for (int linear4 = tid; linear4 < kFloat4PerChunk; linear4 += kCopyThreads) { + const int elem = linear4 * 4; + const int row = elem / kDimK; + const int k = elem - row * kDimK; + float* dst = s_state + (stage * kChunkV + row) * kDimK + k; + const float* src = + state + ((slot * HV + i_hv) * kDimV + v_base + row) * kDimK + k; + cp_async_cg_16b(dst, src); + } + cp_async_commit(); +} + +__device__ __forceinline__ void cp_async_state_chunk(float* s_state, + const float* state, + int slot, int i_hv, int HV, + int chunk) { + cp_async_state_chunk_for(s_state, state, slot, i_hv, HV, chunk); +} + +__device__ __forceinline__ float block_reduce_sum(float value, float* scratch) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + + float warp_total = warp_reduce_sum(value); + if (lane == 0) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +struct Sum2 { + float x; + float y; +}; + +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffffu, x, offset); + y += __shfl_xor_sync(0xffffffffu, y, offset); + } + return {x, y}; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + const float warp_x = warp_reduce_sum(x); + const float warp_y = warp_reduce_sum(y); + if (lane == 0) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +__device__ __forceinline__ Sum2 block_reduce_sum2(float x, float y, + float* scratch) { + return block_reduce_sum2_for(x, y, scratch); +} + +template +__device__ __forceinline__ float block_reduce_sum_active_for(float value, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_total = 0.0f; + if (warp < kReduceWarps) { + warp_total = warp_reduce_sum(value); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_total; + } + __syncthreads(); + + float block_total = 0.0f; + if (warp == 0) { + block_total = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_total = warp_reduce_sum(block_total); + if (lane == 0) { + scratch[0] = block_total; + } + } + __syncthreads(); + return scratch[0]; +} + +template +__device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, + float* scratch) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + + float warp_x = 0.0f; + float warp_y = 0.0f; + if (warp < kReduceWarps) { + warp_x = warp_reduce_sum(x); + warp_y = warp_reduce_sum(y); + } + if (lane == 0 && warp < kReduceWarps) { + scratch[warp] = warp_x; + scratch[kReduceWarps + warp] = warp_y; + } + __syncthreads(); + + float block_x = 0.0f; + float block_y = 0.0f; + if (warp == 0) { + block_x = lane < kReduceWarps ? scratch[lane] : 0.0f; + block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; + block_x = warp_reduce_sum(block_x); + block_y = warp_reduce_sum(block_y); + if (lane == 0) { + scratch[0] = block_x; + scratch[1] = block_y; + } + } + __syncthreads(); + return {scratch[0], scratch[1]}; +} + +template +__global__ +__launch_bounds__(kThreads, 2) void kda_decode_fusion_many_heads_kernel( + const __nv_bfloat16* __restrict__ x_q, + const __nv_bfloat16* __restrict__ x_k, + const __nv_bfloat16* __restrict__ x_v, const float* __restrict__ w_q_t, + const float* __restrict__ w_k_t, const float* __restrict__ w_v_t, + const float* __restrict__ bias_q, const float* __restrict__ bias_k, + const float* __restrict__ bias_v, __nv_bfloat16* __restrict__ cs_q, + __nv_bfloat16* __restrict__ cs_k, __nv_bfloat16* __restrict__ cs_v, + const float* __restrict__ a_log, const __nv_bfloat16* __restrict__ g, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ beta, + const __nv_bfloat16* __restrict__ onorm_g, + const float* __restrict__ onorm_weight, + const int* __restrict__ ssm_state_indices, + const int* __restrict__ cu_seqlens, float* __restrict__ state, + __nv_bfloat16* __restrict__ out, int B, int H, int HV, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides) { + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + int i_n; + int i_hv; + int i_h; + int bos; + int slot; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); +#endif + if constexpr (kUseStaticDecodeLayout) { + if constexpr (kUseHeadGrid) { + i_n = blockIdx.x; + i_hv = blockIdx.y; + } else { + const int nhv = blockIdx.x; + i_n = nhv / kFixedValueHeads; + i_hv = nhv - i_n * kFixedValueHeads; + } + i_h = i_hv; + bos = i_n; + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } else { + const int nhv = blockIdx.x; + i_n = nhv / HV; + i_hv = nhv - i_n * HV; + const int hv_per_h = HV / H; + i_h = i_hv / hv_per_h; + + bos = cu_seqlens == nullptr ? i_n : cu_seqlens[i_n]; + const int eos = cu_seqlens == nullptr ? i_n + 1 : cu_seqlens[i_n + 1]; + if (eos <= bos) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + slot = ssm_state_indices == nullptr ? i_n : ssm_state_indices[i_n]; + } + + constexpr int kLocalDim = kFixedHeads * kDimK; + constexpr int kPackedDim = 3 * kLocalDim; + const int hk_off = i_h * kDimK; + const int hv_off = i_hv * kDimV; + constexpr int hv_count = kFixedValueHeads; + float* const state_for_slot = state + slot * strides.state_slot; + const int64_t conv_slot_offset = slot * strides.conv_slot; + __nv_bfloat16* const cs_q_for_slot = cs_q + conv_slot_offset; + __nv_bfloat16* const cs_k_for_slot = cs_k + conv_slot_offset; + __nv_bfloat16* const cs_v_for_slot = cs_v + conv_slot_offset; + + __shared__ float s_state[2][kChunkV][kDimK]; + __shared__ float s_q[kDimK]; + __shared__ float s_k[kDimK]; + __shared__ float s_decay[kDimK]; + __shared__ float s_v[kDimV]; + __shared__ float s_o[kDimV]; + __shared__ float s_reduce[kThreads]; + __shared__ float s_beta; + float pre_onorm_gate = 0.0f; + float pre_onorm_weight = 0.0f; + + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, 0); + + if constexpr (kUpdateConvState) { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const int64_t xq_idx = bos * strides.x_row + i_h * kDimK + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); + __nv_bfloat16 q_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 q_shift1 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 k_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 q_state = cs_q_for_slot[hk + w * kPackedDim]; + const __nv_bfloat16 k_state = cs_k_for_slot[hk + w * kPackedDim]; + q_acc += __bfloat162float(q_state) * + conv_weight_load(w_q_t, hk, w); + k_acc += __bfloat162float(k_state) * + conv_weight_load(w_k_t, hk, w); + if (w == 1) { + q_shift0 = q_state; + k_shift0 = k_state; + } else if (w == 2) { + q_shift1 = q_state; + k_shift1 = k_state; + } + } + const __nv_bfloat16 q_new = x_q[xq_idx]; + const __nv_bfloat16 k_new = x_k[xq_idx]; + q_acc += __bfloat162float(q_new) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += __bfloat162float(k_new) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + cs_q_for_slot[hk] = q_shift0; + cs_q_for_slot[hk + kPackedDim] = q_shift1; + cs_q_for_slot[hk + 2 * kPackedDim] = q_new; + cs_k_for_slot[hk] = k_shift0; + cs_k_for_slot[hk + kPackedDim] = k_shift1; + cs_k_for_slot[hk + 2 * kPackedDim] = k_new; + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } else { + if (tid < kDimK) { + const int k = tid; + const int hk = hk_off + k; + const float exp_a = + __shfl_sync(0xffffffffu, lane == 0 ? __expf(a_log[i_h]) : 0.0f, 0); + + float q_acc = bias_q == nullptr ? 0.0f : bf16_load(bias_q, hk); + float k_acc = bias_k == nullptr ? 0.0f : bf16_load(bias_k, hk); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hk + w * kPackedDim; + q_acc += bf16_load(cs_q_for_slot, cs_idx) * + conv_weight_load(w_q_t, hk, w); + k_acc += bf16_load(cs_k_for_slot, cs_idx) * + conv_weight_load(w_k_t, hk, w); + } + q_acc += bf16_load(x_q, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_q_t, hk, kKernelWidth - 1); + k_acc += bf16_load(x_k, bos * strides.x_row + i_h * kDimK + k) * + conv_weight_load(w_k_t, hk, kKernelWidth - 1); + + s_q[k] = silu_fast(q_acc); + s_k[k] = silu_fast(k_acc); + + const int64_t gate_idx = bos * kLocalDim + i_hv * kDimK + k; + const float g_raw = bf16_load(g, gate_idx) + dt_bias[hk]; + if constexpr (kUseLowerBound) { + s_decay[k] = __expf(lower_bound * sigmoid_fast(exp_a * g_raw)); + } else { + s_decay[k] = __expf(-exp_a * softplus_fast(g_raw)); + } + } + } + + if constexpr (kUpdateConvState) { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + const int64_t xv_idx = bos * strides.x_row + i_hv * kDimV + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); + __nv_bfloat16 v_shift0 = __float2bfloat16(0.0f); + __nv_bfloat16 v_shift1 = __float2bfloat16(0.0f); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const __nv_bfloat16 v_state = cs_v_for_slot[hvv + w * kPackedDim]; + v_acc += __bfloat162float(v_state) * + conv_weight_load(w_v_t, hvv, w); + if (w == 1) { + v_shift0 = v_state; + } else if (w == 2) { + v_shift1 = v_state; + } + } + const __nv_bfloat16 v_new = x_v[xv_idx]; + v_acc += __bfloat162float(v_new) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + cs_v_for_slot[hvv] = v_shift0; + cs_v_for_slot[hvv + kPackedDim] = v_shift1; + cs_v_for_slot[hvv + 2 * kPackedDim] = v_new; + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } else { + if (tid < kDimV) { + const int v = tid; + const int hvv = hv_off + v; + + float v_acc = bias_v == nullptr ? 0.0f : bf16_load(bias_v, hvv); +#pragma unroll + for (int w = 0; w < kConvStateWidth; ++w) { + const int cs_idx = hvv + w * kPackedDim; + v_acc += bf16_load(cs_v_for_slot, cs_idx) * + conv_weight_load(w_v_t, hvv, w); + } + v_acc += bf16_load(x_v, bos * strides.x_row + i_hv * kDimV + v) * + conv_weight_load(w_v_t, hvv, kKernelWidth - 1); + s_v[v] = silu_fast(v_acc); + + if constexpr (kApplyOnorm && kPreloadOnormParams) { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + v; + pre_onorm_gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + pre_onorm_weight = onorm_weight[v]; + } + } + } + + if (tid == 0) { + const float beta_raw = bf16_load(beta, bos * strides.beta_row + i_hv); + if constexpr (kApplyBetaSigmoid) { + s_beta = sigmoid_fast(beta_raw); + } else { + s_beta = beta_raw; + } + } + __syncthreads(); + + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, hv_count, + 1); + } + + const float q_sq = tid < kDimK ? s_q[tid] * s_q[tid] : 0.0f; + const float k_sq = tid < kDimK ? s_k[tid] * s_k[tid] : 0.0f; + Sum2 qk_sum; + if constexpr (kUseActiveQkReduction) { + qk_sum = block_reduce_sum2_active_for(q_sq, k_sq, s_reduce); + } else { + qk_sum = block_reduce_sum2(q_sq, k_sq, s_reduce); + } + if (tid < kDimK) { + s_q[tid] *= rsqrtf(qk_sum.x + 1.0e-6f) * scale; + s_k[tid] *= rsqrtf(qk_sum.y + 1.0e-6f); + } + __syncthreads(); + + const int k_base = lane * 4; + const float4 q4 = *reinterpret_cast(s_q + k_base); + const float4 k4 = *reinterpret_cast(s_k + k_base); + const float4 decay4 = *reinterpret_cast(s_decay + k_base); + float r_q[4] = {q4.x, q4.y, q4.z, q4.w}; + float r_k[4] = {k4.x, k4.y, k4.z, k4.w}; + float r_decay[4] = {decay4.x, decay4.y, decay4.z, decay4.w}; + float o_sumsq = 0.0f; + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) { + if constexpr (kPrefetchNextStateChunk && kNumChunks > 1) { + if (chunk + 1 < kNumChunks) { + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + } else { + cp_async_wait_all(); + } + if constexpr (!kSkipWarpSync) { + __syncwarp(); + } + + if constexpr (!kPrefetchNextStateChunk) { + if (chunk + 1 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 1); + } + } + +#pragma unroll + for (int row = 0; row < kRowsPerWarp; row += 2) { + const int v_row_a = warp + row * kWarps; + const int v_row_b = warp + (row + 1) * kWarps; + const int v0 = chunk * kChunkV + v_row_a; + const int v1 = chunk * kChunkV + v_row_b; + float h_a_vals[4]; + float h_b_vals[4]; + float dot_hk_a = 0.0f; + float dot_hk_b = 0.0f; + + const float4 raw_h_a = *reinterpret_cast( + &s_state[chunk & 1][v_row_a][k_base]); + const float4 raw_h_b = *reinterpret_cast( + &s_state[chunk & 1][v_row_b][k_base]); + h_a_vals[0] = raw_h_a.x * r_decay[0]; + h_a_vals[1] = raw_h_a.y * r_decay[1]; + h_a_vals[2] = raw_h_a.z * r_decay[2]; + h_a_vals[3] = raw_h_a.w * r_decay[3]; + h_b_vals[0] = raw_h_b.x * r_decay[0]; + h_b_vals[1] = raw_h_b.y * r_decay[1]; + h_b_vals[2] = raw_h_b.z * r_decay[2]; + h_b_vals[3] = raw_h_b.w * r_decay[3]; + dot_hk_a = h_a_vals[0] * r_k[0] + h_a_vals[1] * r_k[1] + + h_a_vals[2] * r_k[2] + h_a_vals[3] * r_k[3]; + dot_hk_b = h_b_vals[0] * r_k[0] + h_b_vals[1] * r_k[1] + + h_b_vals[2] * r_k[2] + h_b_vals[3] * r_k[3]; + + const Sum2 dot_hk = warp_reduce_sum_pair(dot_hk_a, dot_hk_b); + const float v_new0 = (s_v[v0] - dot_hk.x) * s_beta; + const float v_new1 = (s_v[v1] - dot_hk.y) * s_beta; + + float dot_hq_a = 0.0f; + float dot_hq_b = 0.0f; + const int state_idx_a = (i_hv * kDimV + v0) * kDimK + k_base; + const int state_idx_b = (i_hv * kDimV + v1) * kDimK + k_base; + const float h_a_0 = h_a_vals[0] + r_k[0] * v_new0; + const float h_a_1 = h_a_vals[1] + r_k[1] * v_new0; + const float h_a_2 = h_a_vals[2] + r_k[2] * v_new0; + const float h_a_3 = h_a_vals[3] + r_k[3] * v_new0; + const float h_b_0 = h_b_vals[0] + r_k[0] * v_new1; + const float h_b_1 = h_b_vals[1] + r_k[1] * v_new1; + const float h_b_2 = h_b_vals[2] + r_k[2] * v_new1; + const float h_b_3 = h_b_vals[3] + r_k[3] * v_new1; + if constexpr (kComputeOutputBeforeStore) { + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + } else { + store_state_float4( + state_for_slot + state_idx_a, + make_float4(h_a_0, h_a_1, h_a_2, h_a_3)); + store_state_float4( + state_for_slot + state_idx_b, + make_float4(h_b_0, h_b_1, h_b_2, h_b_3)); + dot_hq_a = + h_a_0 * r_q[0] + h_a_1 * r_q[1] + h_a_2 * r_q[2] + h_a_3 * r_q[3]; + dot_hq_b = + h_b_0 * r_q[0] + h_b_1 * r_q[1] + h_b_2 * r_q[2] + h_b_3 * r_q[3]; + } + + const Sum2 dot_hq = warp_reduce_sum_pair(dot_hq_a, dot_hq_b); + if (lane == 0) { + s_o[v0] = dot_hq.x; + s_o[v1] = dot_hq.y; + if constexpr (kApplyOnorm && kAccumulateOnormSumsq) { + o_sumsq += dot_hq.x * dot_hq.x + dot_hq.y * dot_hq.y; + } + } + } + + if constexpr (kPrefetchNextStateChunk) { + if (chunk + 2 < kNumChunks) { + cp_async_state_chunk(&s_state[0][0][0], state_for_slot, 0, i_hv, + hv_count, chunk + 2); + } + } + } + __syncthreads(); + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + if constexpr (kApplyOnorm) { + if constexpr (kAccumulateOnormSumsq) { + if (lane == 0) { + s_reduce[warp] = o_sumsq; + } + __syncthreads(); + + float total_sumsq = 0.0f; + if (warp == 0) { + total_sumsq = lane < kWarps ? s_reduce[lane] : 0.0f; + total_sumsq = warp_reduce_sum(total_sumsq); + if (lane == 0) { + s_reduce[0] = total_sumsq; + } + } + __syncthreads(); + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float raw_o = s_o[tid]; + const float rstd = + rsqrtf(s_reduce[0] / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } else { + const float raw_o = tid < kDimV ? s_o[tid] : 0.0f; + const float o_sq = raw_o * raw_o; + float sumsq; + if constexpr (kUseActiveOnormReduction || kUseActiveQkReduction) { + sumsq = block_reduce_sum_active_for(o_sq, s_reduce); + } else { + sumsq = block_reduce_sum(o_sq, s_reduce); + } + + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + const float rstd = + rsqrtf(sumsq / static_cast(kDimV) + onorm_eps); + float gate; + float weight; + if constexpr (kPreloadOnormParams) { + gate = pre_onorm_gate; + weight = pre_onorm_weight; + } else { + const int64_t gate_idx = i_n * strides.onorm_row + i_hv * kDimV + tid; + gate = sigmoid_fast(bf16_load(onorm_g, gate_idx)); + weight = onorm_weight[tid]; + } + const float y = raw_o * rstd * weight * gate; + out[out_idx] = bf16_store(y); + } + } + } else { + if (tid < kDimV) { + const int64_t out_idx = i_n * kLocalDim + i_hv * kDimV + tid; + out[out_idx] = bf16_store(s_o[tid]); + } + } +} + +template +void launch_kda_decode_many_heads_raw( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, float lower_bound, float scale, + float onorm_eps, KdaDecodeStrides strides, cudaStream_t stream) { + auto kernel = &kda_decode_fusion_many_heads_kernel< + kApplyOnorm, true, kHeads, kHeads, true, false, false, false, false, + false, true, true, true, kUpdateConvState, kUseLowerBound, + kApplyBetaSigmoid>; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(B, kHeads); + config.blockDim = dim3(kThreads); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel, + reinterpret_cast(x_q), + reinterpret_cast(x_k), + reinterpret_cast(x_v), + reinterpret_cast(w_q_t), + reinterpret_cast(w_k_t), + reinterpret_cast(w_v_t), + reinterpret_cast(bias_q), + reinterpret_cast(bias_k), + reinterpret_cast(bias_v), + reinterpret_cast<__nv_bfloat16*>(cs_q), + reinterpret_cast<__nv_bfloat16*>(cs_k), + reinterpret_cast<__nv_bfloat16*>(cs_v), a_log, + reinterpret_cast(g), dt_bias, + reinterpret_cast(beta), + reinterpret_cast(onorm_g), + onorm_weight, ssm_state_indices, cu_seqlens, state, + reinterpret_cast<__nv_bfloat16*>(out), B, H, HV, + lower_bound, scale, onorm_eps, strides); +} + +template +void launch_kda_decode_many_heads_selected( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool update_conv_cache, float lower_bound, + float scale, float onorm_eps, KdaDecodeStrides strides, + cudaStream_t stream) { +#define LAUNCH_KDA_DECODE(NUM_HEADS) \ + do { \ + if (update_conv_cache) { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } else { \ + launch_kda_decode_many_heads_raw( \ + x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, \ + cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, \ + ssm_state_indices, cu_seqlens, state, out, B, H, HV, lower_bound, \ + scale, onorm_eps, strides, stream); \ + } \ + } while (false) + + switch (H) { + case 12: + LAUNCH_KDA_DECODE(12); + break; + case 24: + LAUNCH_KDA_DECODE(24); + break; + case 48: + LAUNCH_KDA_DECODE(48); + break; + case 96: + LAUNCH_KDA_DECODE(96); + break; + default: + STD_TORCH_CHECK(false, "Unsupported number of heads: ", H); + } + +#undef LAUNCH_KDA_DECODE +} + +struct KdaDecodeLaunchParams { + const void* x_q; + const void* x_k; + const void* x_v; + const void* w_q_t; + const void* w_k_t; + const void* w_v_t; + const void* bias_q; + const void* bias_k; + const void* bias_v; + void* cs_q; + void* cs_k; + void* cs_v; + const float* a_log; + const void* g; + const float* dt_bias; + const void* beta; + const void* onorm_g; + const float* onorm_weight; + const int* ssm_state_indices; + const int* cu_seqlens; + float* state; + void* out; + int B; + int H; + int HV; + bool update_conv_cache; + float lower_bound; + float scale; + float onorm_eps; + KdaDecodeStrides strides; + cudaStream_t stream; +}; + +template +void launch_kda_decode_selected_backend(const KdaDecodeLaunchParams& p) { + launch_kda_decode_many_heads_selected( + p.x_q, p.x_k, p.x_v, p.w_q_t, p.w_k_t, p.w_v_t, p.bias_q, p.bias_k, + p.bias_v, p.cs_q, p.cs_k, p.cs_v, p.a_log, p.g, p.dt_bias, p.beta, + p.onorm_g, p.onorm_weight, p.ssm_state_indices, p.cu_seqlens, p.state, + p.out, p.B, p.H, p.HV, p.update_conv_cache, p.lower_bound, p.scale, + p.onorm_eps, p.strides, p.stream); +} + +template +void dispatch_kda_decode_beta(const KdaDecodeLaunchParams& p, + bool apply_beta_sigmoid) { + if (apply_beta_sigmoid) { + launch_kda_decode_selected_backend(p); + } else { + launch_kda_decode_selected_backend(p); + } +} + +template +void dispatch_kda_decode_decay(const KdaDecodeLaunchParams& p, + bool use_lower_bound, bool apply_beta_sigmoid) { + if (use_lower_bound) { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } else { + dispatch_kda_decode_beta(p, apply_beta_sigmoid); + } +} + +void dispatch_kda_decode_features(const KdaDecodeLaunchParams& p, + bool apply_onorm, bool use_lower_bound, + bool apply_beta_sigmoid) { + if (apply_onorm) { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } else { + dispatch_kda_decode_decay(p, use_lower_bound, apply_beta_sigmoid); + } +} + +} // namespace + +extern "C" void launch_kda_decode_many_heads_cuda( + const void* x_q, const void* x_k, const void* x_v, const void* w_q_t, + const void* w_k_t, const void* w_v_t, const void* bias_q, + const void* bias_k, const void* bias_v, void* cs_q, void* cs_k, void* cs_v, + const float* a_log, const void* g, const float* dt_bias, const void* beta, + const void* onorm_g, const float* onorm_weight, + const int* ssm_state_indices, const int* cu_seqlens, float* state, + void* out, int B, int H, int HV, bool apply_onorm, bool update_conv_cache, + bool use_lower_bound, bool apply_beta_sigmoid, float lower_bound, + float scale, float onorm_eps, const int64_t* raw_strides, + cudaStream_t stream) { + const KdaDecodeStrides strides{raw_strides[0], raw_strides[1], raw_strides[2], + raw_strides[3], raw_strides[4]}; + const KdaDecodeLaunchParams params{x_q, + x_k, + x_v, + w_q_t, + w_k_t, + w_v_t, + bias_q, + bias_k, + bias_v, + cs_q, + cs_k, + cs_v, + a_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + ssm_state_indices, + cu_seqlens, + state, + out, + B, + H, + HV, + update_conv_cache, + lower_bound, + scale, + onorm_eps, + strides, + stream}; + dispatch_kda_decode_features(params, apply_onorm, use_lower_bound, + apply_beta_sigmoid); +} + +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps) { + using torch::headeronly::ScalarType; + constexpr int kHeadDim = 128; + constexpr int kConvWidth = 4; + + STD_TORCH_CHECK(x.is_cuda() && x.scalar_type() == ScalarType::BFloat16, + "x must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == ScalarType::Float, + "weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + conv_state.is_cuda() && conv_state.scalar_type() == ScalarType::BFloat16, + "conv_state must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_g.is_cuda() && raw_g.scalar_type() == ScalarType::BFloat16, + "raw_g must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + raw_beta.is_cuda() && raw_beta.scalar_type() == ScalarType::BFloat16, + "raw_beta must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, + "A_log must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + dt_bias.is_cuda() && dt_bias.scalar_type() == ScalarType::Float, + "dt_bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(state.is_cuda() && state.scalar_type() == ScalarType::Float, + "state must be a CUDA float32 tensor"); + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, + "out must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK( + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, + "state_indices must be a CUDA int32 tensor"); + + STD_TORCH_CHECK(x.dim() == 2, "x must have shape [B, 3 * H * 128]"); + int const batch_size = static_cast(x.size(0)); + int64_t const qkv_width = x.size(1); + STD_TORCH_CHECK(qkv_width % (3 * kHeadDim) == 0, + "x must have shape [B, 3 * H * 128]"); + int64_t const num_heads = qkv_width / (3 * kHeadDim); + STD_TORCH_CHECK( + num_heads == 12 || num_heads == 24 || num_heads == 48 || num_heads == 96, + "H must be 12, 24, 48, or 96, got ", num_heads); + STD_TORCH_CHECK(batch_size > 0, + "KDA decode fusion requires at least one row"); + int const dim = num_heads * kHeadDim; + + STD_TORCH_CHECK(weight.dim() == 3 && weight.is_contiguous() && + weight.size(0) == 3 && weight.size(1) == kConvWidth && + weight.size(2) == dim, + "weight must have shape [3, 4, H * 128]"); + STD_TORCH_CHECK(conv_state.dim() == 3 && conv_state.size(1) == 3 * dim && + conv_state.size(2) == kConvWidth - 1, + "conv_state must have shape [slots, 3 * H * 128, 3]"); + STD_TORCH_CHECK(raw_g.dim() == 4 && raw_g.size(0) == 1 && + raw_g.size(1) == batch_size && + raw_g.size(2) == num_heads && raw_g.size(3) == kHeadDim, + "raw_g must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(raw_beta.dim() == 3 && raw_beta.size(0) == 1 && + raw_beta.size(1) == batch_size && + raw_beta.size(2) == num_heads, + "raw_beta must have shape [1, B, H]"); + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_heads, + "A_log must be contiguous with H elements"); + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == dim, + "dt_bias must be contiguous with H * 128 elements"); + STD_TORCH_CHECK( + state_indices.is_contiguous() && state_indices.numel() == batch_size, + "state_indices must be contiguous with B elements"); + STD_TORCH_CHECK(state.dim() == 4 && state.size(1) == num_heads && + state.size(2) == kHeadDim && state.size(3) == kHeadDim, + "state must have shape [slots, H, 128, 128]"); + STD_TORCH_CHECK(out.dim() == 4 && out.size(0) == 1 && + out.size(1) == batch_size && out.size(2) == num_heads && + out.size(3) == kHeadDim, + "out must have shape [1, B, H, 128]"); + STD_TORCH_CHECK(x.stride(1) == 1, + "x must be contiguous in its channel dimension"); + STD_TORCH_CHECK(conv_state.stride(0) >= 3 * dim * (kConvWidth - 1) && + conv_state.stride(1) == 1 && + conv_state.stride(2) == 3 * dim, + "conv_state must use the SD cache layout"); + STD_TORCH_CHECK(state.stride(0) >= num_heads * kHeadDim * kHeadDim && + state.stride(1) == kHeadDim * kHeadDim && + state.stride(2) == kHeadDim && state.stride(3) == 1, + "state must have contiguous [H, 128, 128] slot contents"); + STD_TORCH_CHECK(raw_g.is_contiguous(), "raw_g must be contiguous"); + STD_TORCH_CHECK(raw_beta.stride(2) == 1, + "raw_beta must be contiguous in its head dimension"); + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + + bool const apply_onorm = output_gate.has_value(); + STD_TORCH_CHECK(apply_onorm == norm_weight.has_value(), + "output_gate and norm_weight must be provided together"); + void const* output_gate_ptr = nullptr; + float const* norm_weight_ptr = nullptr; + int64_t output_gate_row_stride = 0; + if (apply_onorm) { + STD_TORCH_CHECK(output_gate->is_cuda() && + output_gate->scalar_type() == ScalarType::BFloat16, + "output_gate must be a CUDA bfloat16 tensor"); + bool const gate_is_3d = + output_gate->dim() == 3 && output_gate->size(0) == batch_size && + output_gate->size(1) == num_heads && output_gate->size(2) == kHeadDim; + bool const gate_is_4d = + output_gate->dim() == 4 && output_gate->size(0) == 1 && + output_gate->size(1) == batch_size && + output_gate->size(2) == num_heads && output_gate->size(3) == kHeadDim; + STD_TORCH_CHECK(gate_is_3d || gate_is_4d, + "output_gate must have shape [B, H, 128] or " + "[1, B, H, 128]"); + int const row_dim = gate_is_3d ? 0 : 1; + STD_TORCH_CHECK(output_gate->stride(output_gate->dim() - 1) == 1, + "output_gate must be contiguous in its last dimension"); + STD_TORCH_CHECK(output_gate->stride(row_dim + 1) == kHeadDim, + "output_gate must have contiguous head rows"); + STD_TORCH_CHECK(norm_weight->is_cuda() && + norm_weight->scalar_type() == ScalarType::Float, + "norm_weight must be a CUDA float32 tensor"); + STD_TORCH_CHECK( + norm_weight->is_contiguous() && norm_weight->numel() == kHeadDim, + "norm_weight must be contiguous with 128 elements"); + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); + output_gate_ptr = output_gate->data_ptr(); + norm_weight_ptr = static_cast(norm_weight->data_ptr()); + output_gate_row_stride = output_gate->stride(row_dim); + } + + void const* bias_ptr = nullptr; + if (bias.has_value()) { + STD_TORCH_CHECK(bias->is_cuda() && bias->scalar_type() == ScalarType::Float, + "bias must be a CUDA float32 tensor"); + STD_TORCH_CHECK(bias->is_contiguous() && bias->numel() == 3 * dim, + "bias must be contiguous with 3 * H * 128 elements"); + bias_ptr = bias->data_ptr(); + } + + auto const* x_ptr = static_cast(x.data_ptr()); + auto const* weight_ptr = static_cast(weight.data_ptr()); + auto* conv_ptr = static_cast(conv_state.data_ptr()); + auto const* bias_bytes = static_cast(bias_ptr); + int64_t const segment_bytes = dim * sizeof(__nv_bfloat16); + int64_t const weight_segment_bytes = + dim * kConvWidth * static_cast(sizeof(float)); + int64_t const conv_segment_bytes = + dim * conv_state.stride(1) * sizeof(__nv_bfloat16); + int64_t const bias_segment_bytes = dim * sizeof(float); + std::array const strides{ + x.stride(0), raw_beta.stride(1), output_gate_row_stride, + conv_state.stride(0), state.stride(0), + }; + bool const use_lower_bound = lower_bound.has_value(); + float const lower_bound_value = + use_lower_bound ? static_cast(*lower_bound) : 0.0f; + + torch::stable::accelerator::DeviceGuard const device_guard( + x.get_device_index()); + cudaStream_t const stream = get_current_cuda_stream(x.get_device_index()); + launch_kda_decode_many_heads_cuda( + x_ptr, x_ptr + segment_bytes, x_ptr + 2 * segment_bytes, weight_ptr, + weight_ptr + weight_segment_bytes, weight_ptr + 2 * weight_segment_bytes, + bias_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + bias_segment_bytes, + bias_bytes == nullptr ? nullptr : bias_bytes + 2 * bias_segment_bytes, + conv_ptr, conv_ptr + conv_segment_bytes, + conv_ptr + 2 * conv_segment_bytes, + static_cast(a_log.data_ptr()), raw_g.data_ptr(), + static_cast(dt_bias.data_ptr()), raw_beta.data_ptr(), + output_gate_ptr, norm_weight_ptr, + static_cast(state_indices.data_ptr()), nullptr, + static_cast(state.data_ptr()), out.data_ptr(), batch_size, + num_heads, num_heads, apply_onorm, true, use_lower_bound, true, + lower_bound_value, 0.08838834764831845f, static_cast(norm_eps), + strides.data(), stream); + cudaError_t const error = cudaGetLastError(); + STD_TORCH_CHECK( + error == cudaSuccess, + "Kimi K3 KDA decode kernel launch failed: ", cudaGetErrorString(error)); +} diff --git a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu index da9ef44d03f..e329e067b26 100644 --- a/csrc/libtorch_stable/moe/grouped_topk_kernels.cu +++ b/csrc/libtorch_stable/moe/grouped_topk_kernels.cu @@ -25,6 +25,7 @@ #include "libtorch_stable/torch_utils.h" #include +#include #include #include #include @@ -448,7 +449,8 @@ enum ScoringFunc { SCORING_SIGMOID = 1 // apply sigmoid }; -// Efficient sigmoid approximation from TensorRT-LLM +// Adapted from +// https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu __device__ inline float sigmoid_accurate(float x) { return 0.5f * tanhf(0.5f * x) + 0.5f; } @@ -890,6 +892,434 @@ __global__ void grouped_topk_fused_small_expert_count_kernel( #endif } +// Adapted from +// https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh +namespace single_group_topk { +namespace detail { + +static constexpr int BlockDim = 256; +static constexpr uint32_t FullWarpMask = 0xffffffffU; +static constexpr float InvalidScore = -INFINITY; + +// TopK-only tuning: use wider workers and keep these tiers on the block path. +template +static constexpr bool UseTunedBlockPath = + MaxNumTopExperts == 16 && (MaxNumExperts == 896 || MaxNumExperts == 1024); + +template +__device__ __forceinline__ void preprocess_score(T input, BiasT correction_bias, + float& unbiased_score, + float& selection_score) { + unbiased_score = 0.0F; + selection_score = InvalidScore; + float const input_float = cuda_cast(input); + float const bias = cuda_cast(correction_bias); + if (!is_finite(input_float) || !is_finite(bias)) { + return; + } + + float const unbiased = apply_scoring(input_float); + float const biased = unbiased + bias; + if constexpr (SF == SCORING_NONE) { + if (!is_finite(biased)) { + return; + } + } + unbiased_score = unbiased; + selection_score = biased == 0.0F ? 0.0F : biased; +} + +template +__device__ __forceinline__ void write_outputs( + cg::thread_block_tile const& warp, float lane_selection_score, + float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token, + int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize, + float routed_scaling_factor) { + bool const finite_selection = + lane < topk && lane_selection_score != InvalidScore; + lane_unbiased = finite_selection ? lane_unbiased : 0.0F; + unsigned const finite_mask = __ballot_sync(FullWarpMask, finite_selection); + float const sum = cg::reduce(warp, lane_unbiased, cg::plus{}); + + if (lane < topk) { + float output = 0.0F; + if (finite_mask == 0) { + if (renormalize) { + output = 1.0F / static_cast(topk); + } + } else if (finite_selection) { + float scale = routed_scaling_factor; + if (renormalize) { + scale /= sum + 1e-20F; + } + output = lane_unbiased * scale; + } + + int64_t const output_index = int64_t{token} * topk + lane; + topk_values[output_index] = output; + topk_indices[output_index] = static_cast(lane_expert); + } +} + +template +__global__ void __launch_bounds__(BlockDim) + single_group_topk_block_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_experts, int64_t topk, + bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WorkerValuesPerLane = + UseTunedBlockPath ? 8 : 4; + static constexpr int ExpertsPerWorkerWarp = WorkerValuesPerLane * WARP_SIZE; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + static constexpr int NumWorkerWarps = + (MaxNumExperts + ExpertsPerWorkerWarp - 1) / ExpertsPerWorkerWarp; + static constexpr int NumIntermediate = NumWorkerWarps * MaxNumTopExperts; + static constexpr int MergeValuesPerLane = + (NumIntermediate + WARP_SIZE - 1) / WARP_SIZE; + static constexpr bool LaneOwnedResourcesFit = + NumWorkerWarps <= BlockDim / WARP_SIZE && MergeValuesPerLane <= 64; + static constexpr bool UseHierarchicalLaneTopK = + LaneOwnedRange::kEnabled && LaneOwnedResourcesFit; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + __shared__ float __attribute((aligned(128))) biased_scores[MaxNumExperts]; + __shared__ float __attribute((aligned(128))) unbiased_scores[MaxNumExperts]; + + int32_t const token = static_cast(blockIdx.x); + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + T const* token_scores = scores + int64_t{token} * num_experts; + + for (int32_t expert = static_cast(threadIdx.x); + expert < num_experts_i32; expert += BlockDim) { + preprocess_score(token_scores[expert], bias[expert], + unbiased_scores[expert], + biased_scores[expert]); + } + __syncthreads(); + + auto warp = cg::tiled_partition(cg::this_thread_block()); + + if constexpr (UseHierarchicalLaneTopK) { + __shared__ float + __attribute((aligned(128))) intermediate_scores[NumIntermediate]; + __shared__ int32_t + __attribute((aligned(128))) intermediate_indices[NumIntermediate]; + + if (warp_id < NumWorkerWarps) { + float local_scores[WorkerValuesPerLane]; + int32_t local_indices[WorkerValuesPerLane]; +#pragma unroll + for (int index = 0; index < WorkerValuesPerLane; ++index) { + int32_t const expert = + warp_id * ExpertsPerWorkerWarp + index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + if (lane < MaxNumTopExperts) { + int32_t const intermediate = warp_id * MaxNumTopExperts + lane; + bool const active = lane < topk_i32; + intermediate_scores[intermediate] = active ? lane_score : InvalidScore; + intermediate_indices[intermediate] = + active ? lane_expert : MaxNumExperts; + } + } + __syncthreads(); + + if (warp_id != 0) { + return; + } + + float merge_scores[MergeValuesPerLane]; + int32_t merge_indices[MergeValuesPerLane]; +#pragma unroll + for (int index = 0; index < MergeValuesPerLane; ++index) { + int32_t const intermediate = index * WARP_SIZE + lane; + bool const active = intermediate < NumIntermediate; + merge_scores[index] = + active ? intermediate_scores[intermediate] : InvalidScore; + merge_indices[index] = + active ? intermediate_indices[intermediate] : MaxNumExperts; + } + + float lane_score; + int32_t lane_expert; + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, merge_scores, merge_indices, + InvalidScore, lane); + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } else { + if (warp_id != 0) { + return; + } + + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + local_scores[index] = + expert < num_experts_i32 ? biased_scores[expert] : InvalidScore; + local_indices[index] = expert; + } + + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + float const lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + int32_t const lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + float const lane_unbiased = + lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32 + ? unbiased_scores[lane_expert] + : 0.0F; + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct WarpTopKLaunchConfig { + static constexpr int DefaultBlockDim = + MaxNumExperts <= 1024 ? MaxNumExperts : 1024; + static constexpr int BlockDim = DefaultBlockDim > 256 ? 256 : DefaultBlockDim; + static constexpr int NumWarps = BlockDim / WARP_SIZE; + static constexpr int MaxBlockScale = + (DefaultBlockDim + BlockDim - 1) / BlockDim; + static constexpr int MaxBlocks = 1024 * MaxBlockScale; + + static_assert(BlockDim % WARP_SIZE == 0); + + static uint32_t grid_dim(int64_t num_tokens) { + int64_t const token_blocks = (num_tokens + NumWarps - 1) / NumWarps; + int64_t const selected = + token_blocks < MaxBlocks ? token_blocks : MaxBlocks; + return static_cast(selected > 0 ? selected : 1); + } +}; + +template +__global__ void __launch_bounds__(WarpTopKLaunchConfig::BlockDim) + single_group_topk_warp_kernel(T const* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, + int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, + float routed_scaling_factor, + bool enable_pdl) { + static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; + static constexpr int WarpBlockDim = + WarpTopKLaunchConfig::BlockDim; + using LaneOwnedRange = + reduce_topk::HighExpertLaneOwnedTopKRange; + + static_assert(NumChunks <= 64); + static_assert(MaxNumTopExperts <= WARP_SIZE); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaGridDependencySynchronize(); + } +#endif + + int32_t const lane = static_cast(threadIdx.x) % WARP_SIZE; + int32_t const warp_id = static_cast(threadIdx.x) / WARP_SIZE; + int32_t const global_warp = + static_cast(blockIdx.x) * WarpBlockDim / WARP_SIZE + warp_id; + int32_t const global_warp_stride = + static_cast(gridDim.x) * WarpBlockDim / WARP_SIZE; + int32_t const num_experts_i32 = static_cast(num_experts); + int32_t const topk_i32 = static_cast(topk); + auto warp = cg::tiled_partition(cg::this_thread_block()); + + for (int32_t token = global_warp; token < num_tokens; + token += global_warp_stride) { + T const* token_scores = scores + int64_t{token} * num_experts; + float local_scores[NumChunks]; + int32_t local_indices[NumChunks]; +#pragma unroll + for (int index = 0; index < NumChunks; ++index) { + int32_t const expert = index * WARP_SIZE + lane; + float unbiased; + float selection; + if (expert < num_experts_i32) { + preprocess_score(token_scores[expert], bias[expert], + unbiased, selection); + } else { + selection = InvalidScore; + } + local_scores[index] = selection; + local_indices[index] = expert; + } + + float lane_score; + int32_t lane_expert; + if constexpr (LaneOwnedRange::kEnabled) { + reduce_topk::reduceTopKForLane( + warp, lane_score, lane_expert, local_scores, local_indices, + InvalidScore, lane); + } else { + float top_scores[MaxNumTopExperts]; + int32_t top_experts[MaxNumTopExperts]; + reduce_topk::reduceTopK(warp, top_scores, top_experts, local_scores, + local_indices, InvalidScore, topk_i32); + lane_score = lane < topk_i32 ? top_scores[lane] : InvalidScore; + lane_expert = lane < topk_i32 ? top_experts[lane] : -1; + } + + float lane_unbiased = 0.0F; + if (lane < topk_i32 && lane_expert >= 0 && lane_expert < num_experts_i32) { + lane_unbiased = lane_score - cuda_cast(bias[lane_expert]); + } + write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, + topk_i32, topk_values, topk_indices, renormalize, + routed_scaling_factor); + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (enable_pdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +struct Tier { + static constexpr int kExperts = Experts; + static constexpr int kTopK = TopK; +}; + +template +struct TierList {}; + +using SigmoidBiasTiers = + TierList, Tier<256, 8>, Tier<384, 8>, Tier<512, 8>, + Tier<512, 22>, Tier<768, 16>, Tier<896, 16>, Tier<1024, 16>>; + +using PrecomputedSoftmaxBiasTiers = + TierList, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, + Tier<256, 16>, Tier<512, 8>, Tier<512, 16>, Tier<512, 22>, + Tier<512, 32>, Tier<576, 8>, Tier<768, 16>, Tier<896, 16>, + Tier<1024, 16>>; + +template +void launch(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + config.dynamicSmemBytes = 0; + bool const use_block_kernel = + UseTunedBlockPath || + MaxNumExperts > 1024 || num_experts >= 1024 || + (num_experts >= 256 && num_tokens <= 1024); + if (use_block_kernel) { + config.gridDim = static_cast(num_tokens); + config.blockDim = BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_block_kernel, + scores, topk_values, topk_indices, bias, num_experts, topk, renormalize, + static_cast(routed_scaling_factor), enable_pdl); + } else { + using WarpConfig = WarpTopKLaunchConfig; + config.gridDim = WarpConfig::grid_dim(num_tokens); + config.blockDim = WarpConfig::BlockDim; + cudaLaunchKernelEx( + &config, + &single_group_topk_warp_kernel, + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, static_cast(routed_scaling_factor), enable_pdl); + } +} + +template +bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t, + int64_t, bool, double, bool, cudaLaunchConfig_t&) { + return false; +} + +template +bool dispatch(TierList*, T* scores, float* topk_values, + IdxT* topk_indices, BiasT const* bias, int64_t num_tokens, + int64_t num_experts, int64_t topk, bool renormalize, + double routed_scaling_factor, bool enable_pdl, + cudaLaunchConfig_t& config) { + if (num_experts <= First::kExperts && topk <= First::kTopK) { + launch( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + return true; + } + return dispatch( + static_cast*>(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); +} + +} // namespace detail + +template +bool invoke(T* scores, float* topk_values, IdxT* topk_indices, + BiasT const* bias, int64_t num_tokens, int64_t num_experts, + int64_t topk, bool renormalize, double routed_scaling_factor, + bool enable_pdl, cudaLaunchConfig_t& config) { + static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID); + if constexpr (SF == SCORING_SIGMOID) { + return detail::dispatch( + static_cast(nullptr), scores, topk_values, + topk_indices, bias, num_tokens, num_experts, topk, renormalize, + routed_scaling_factor, enable_pdl, config); + } else { + return detail::dispatch( + static_cast(nullptr), scores, + topk_values, topk_indices, bias, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); + } +} + +} // namespace single_group_topk + template void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, int64_t const num_tokens, @@ -905,6 +1335,12 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; config.numAttrs = 1; config.attrs = attrs; + if (n_group == 1 && topk_group == 1 && + single_group_topk::invoke( + scores, topk_values, topk_indices, bias, num_tokens, num_experts, + topk, renormalize, routed_scaling_factor, enable_pdl, config)) { + return; + } // Check if we can use the optimized // grouped_topk_fused_small_expert_count_kernel diff --git a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh index 70e21cf8773..6eadae0def8 100644 --- a/csrc/libtorch_stable/moe/moeTopKFuncs.cuh +++ b/csrc/libtorch_stable/moe/moeTopKFuncs.cuh @@ -1,6 +1,7 @@ /* * Adapted from * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + * https://github.com/flashinfer-ai/flashinfer/blob/06400d062a2d51564bbe781f6f811d0b75ca593e/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh * Copyright (c) 2026, The vLLM team. * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights * reserved. SPDX-License-Identifier: Apache-2.0 @@ -23,6 +24,9 @@ #include #include +#include +#include + namespace vllm { namespace moe { namespace reduce_topk { @@ -38,11 +42,10 @@ struct TopKRedType { "Top K reduction only implemented for int, float, float16 and bfloat16"); using TypeCmp = std::conditional_t; - using IdxT = std::conditional_t; static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; static constexpr int kMaxIdx = 65535; - TypeCmp compValIdx; + TypeCmp compVal; static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { auto valueBits = cub::Traits::TwiddleIn( @@ -69,69 +72,175 @@ struct TopKRedType { __host__ __device__ TopKRedType() = default; __host__ __device__ TopKRedType(T val, int32_t idx) - : compValIdx(makeCmpVal(val, idx)) {} + : compVal(makeCmpVal(val, idx)) {} - __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + __host__ __device__ operator TypeCmp() const noexcept { return compVal; } __device__ inline TypeCmp reduce( cg::thread_block_tile const& warp) { - return cg::reduce(warp, compValIdx, cg::greater{}); +#ifdef __CUDA_ARCH__ + static constexpr bool kHAS_FAST_REDUX = (__CUDA_ARCH__ / 100) >= 10; +#else + static constexpr bool kHAS_FAST_REDUX = false; +#endif + if constexpr (!kHAS_FAST_REDUX) { + return cg::reduce(warp, compVal, cg::greater{}); + } else if constexpr (sizeof(TypeCmp) == 8) { + uint32_t hi = static_cast(compVal >> 32); + uint32_t lo = static_cast(compVal & 0xffffffffu); + uint32_t maxHi; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxHi) + : "r"(hi)); + uint32_t loContrib = hi == maxHi ? lo : 0u; + uint32_t maxLo; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(maxLo) + : "r"(loContrib)); + return (static_cast(maxHi) << 32) | static_cast(maxLo); + } else { + TypeCmp result; + asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(result) + : "r"(compVal)); + return result; + } } }; -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct TopKIdx { - // by default, empty +template +struct IsPowerOf2 { + static constexpr bool value = N > 0 && (N & (N - 1)) == 0; }; -template -struct TopKIdx { - static constexpr int K = K_; - int32_t val[K]; +template +struct NextPow2 { + private: + static constexpr unsigned u = static_cast(N - 1); + static constexpr unsigned s1 = u | (u >> 1); + static constexpr unsigned s2 = s1 | (s1 >> 2); + static constexpr unsigned s3 = s2 | (s2 >> 4); + static constexpr unsigned s4 = s3 | (s3 >> 8); + static constexpr unsigned s5 = s4 | (s4 >> 16); + + public: + static constexpr int value = N <= 1 ? 1 : static_cast(s5 + 1); }; -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#define TOPK_SWAP(I, J) \ - { \ - auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \ - auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \ - topK[I].compValIdx = pairMax; \ - topK[J].compValIdx = pairMin; \ +template +__device__ __forceinline__ void topkCompareSwap(T* a) { + if constexpr (A < Size && B < Size) { + if (a[A] < a[B]) { + T tmp = a[A]; + a[A] = a[B]; + a[B] = tmp; + } + } else { + (void)a; } +} + +template +__device__ __forceinline__ void topkMergePairs(T* a) { + if constexpr (I + Step < End) { + topkCompareSwap(a); + topkMergePairs(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkOEM(T* a) { + constexpr int M = R * 2; + if constexpr (M < N) { + topkOEM(a); + topkOEM(a); + topkMergePairs(a); + } else if constexpr (R < N) { + topkCompareSwap(a); + } else { + (void)a; + } +} + +template +__device__ __forceinline__ void topkSortBatcher(T* a) { + if constexpr (N > 1) { + constexpr int Half = N / 2; + topkSortBatcher(a); + topkSortBatcher(a); + topkOEM(a); + } else { + (void)a; + } +} template -struct Sort; +struct Sort { + static_assert(N > 0 && N <= 64, "Sort only supports N in range [1, 64]"); + + static __device__ void run(RedType* topK) { + if constexpr (IsPowerOf2::value) { +#pragma unroll + for (int k = 2; k <= N; k *= 2) { +#pragma unroll + for (int j = k / 2; j > 0; j /= 2) { +#pragma unroll + for (int i = 0; i < N; ++i) { + int ixj = i ^ j; + if (ixj > i) { + if ((i & k) == 0) { + if (topK[i].compVal < topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } else { + if (topK[i].compVal > topK[ixj].compVal) { + auto tmp = topK[i].compVal; + topK[i].compVal = topK[ixj].compVal; + topK[ixj].compVal = tmp; + } + } + } + } + } + } + } else { + constexpr int P = NextPow2::value; + topkSortBatcher<0, P, N, RedType>(topK); + } + } +}; template struct Sort<1, RedType> { - static __device__ void run(RedType* topK) {} + static __device__ void run(RedType*) {} }; template struct Sort<2, RedType> { - static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } + static __device__ void run(RedType* topK) { topkCompareSwap<0, 1, 2>(topK); } }; template struct Sort<3, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 1); - TOPK_SWAP(1, 2); - TOPK_SWAP(0, 1); + topkCompareSwap<0, 1, 3>(topK); + topkCompareSwap<1, 2, 3>(topK); + topkCompareSwap<0, 1, 3>(topK); } }; template struct Sort<4, RedType> { static __device__ void run(RedType* topK) { - TOPK_SWAP(0, 2); - TOPK_SWAP(1, 3); - TOPK_SWAP(0, 1); - TOPK_SWAP(2, 3); - TOPK_SWAP(1, 2); + topkCompareSwap<0, 2, 4>(topK); + topkCompareSwap<1, 3, 4>(topK); + topkCompareSwap<0, 1, 4>(topK); + topkCompareSwap<2, 3, 4>(topK); + topkCompareSwap<1, 2, 4>(topK); } }; @@ -147,110 +256,112 @@ __forceinline__ __device__ void reduceTopK( typename RedType::TypeCmp packedMax{}; #pragma unroll for (int kk = 0; kk < actualK; ++kk) { - topK = - kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK; - // get the next largest value + topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK; packedMax = topK.reduce(warp); RedType::unpack(out[kk], outIdx[kk], packedMax); } }; -template -__device__ void reduceTopKFunc(cg::thread_block_tile const& warp, - Type (&out)[K], int32_t (&outIdx)[K], - Type (&value)[N], int32_t (&idx)[N], - Type minValue, int actualK = K) { - static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); - static_assert(N > 0, "Top K must have N > 0"); - static_assert(N < 5, - "Only support candidates number less than or equal to 128"); - using RedType = TopKRedType; - RedType topK[N]; -#pragma unroll - for (int nn = 0; nn < N; ++nn) { - topK[nn] = RedType{value[nn], idx[nn]}; - } - - if constexpr (!IsSorted) { - Sort::run(topK); - } - typename RedType::TypeCmp packedMax{}; -#pragma unroll - for (int kk = 0; kk < actualK; ++kk) { - bool update = kk > 0 && packedMax == topK[0].compValIdx; -#pragma unroll - for (int nn = 0; nn < N; ++nn) { - topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} - : update ? topK[nn + 1] - : topK[nn]; - } - // get the next largest value - packedMax = topK[0].reduce(warp); - RedType::unpack(out[kk], outIdx[kk], packedMax); - } -}; - template __forceinline__ __device__ void reduceTopK( cg::thread_block_tile const& warp, Type (&out)[K], int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], Type const minValue, int actualK = K) { static_assert(K > 0, "Top K must have K > 0"); - static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); static_assert(N > 0, "Top K must have N > 0"); - static_assert( - N <= 16, - "Only support candidates number less than or equal to 16*32=512"); - static_assert(N <= 4 || N % 4 == 0, - "Only support candidates number is a multiple of 4*32=128 or " - "less than or equal to 4"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } - if constexpr (N <= 4) { - reduceTopKFunc(warp, out, outIdx, value, idx, minValue, - actualK); - } else { - constexpr int numLoops = N / 4; - constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; + Sort::run(topK); - Type topKBufferValue[numResults]; - int32_t topKBufferIdx[numResults]; - int32_t laneIdx = threadIdx.x % kWARP_SIZE; - - for (int ii = 0; ii < numResults; ++ii) { - topKBufferValue[ii] = minValue; - topKBufferIdx[ii] = ii * kWARP_SIZE - 1; + typename RedType::TypeCmp packedMax{}; + for (int kk = 0; kk < actualK; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compVal; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; } - for (int loop = 0; loop < numLoops; ++loop) { - int start = loop * 4; - Type topKValue[K]; - int32_t topKIdx[K]; - Type inValue[4]; - int32_t inIdx[4]; - for (int i = 0; i < 4; ++i) { - inValue[i] = value[start + i]; - inIdx[i] = idx[start + i]; - } - reduceTopKFunc(warp, topKValue, topKIdx, inValue, inIdx, - minValue, actualK); - int inOffset = laneIdx % K; - if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) { - topKBufferValue[0] = topKValue[inOffset]; - topKBufferIdx[0] = topKIdx[inOffset]; - } - if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) { - topKBufferValue[1] = topKValue[inOffset]; - topKBufferIdx[1] = topKIdx[inOffset]; - } - } - - reduceTopKFunc(warp, out, outIdx, topKBufferValue, - topKBufferIdx, minValue, actualK); + packedMax = topK[0].reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); } }; -#undef TOPK_SWAP +template +struct LaneOwnedTopKRange { + static_assert(MinExperts > 0 && MinExperts <= MaxExperts); + static_assert(MinTopExperts > 0 && MinTopExperts <= MaxTopExperts); + static constexpr bool kEnabled = + NumExperts >= MinExperts && NumExperts <= MaxExperts && + NumTopExperts >= MinTopExperts && NumTopExperts <= MaxTopExperts; +}; + +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_EXPERTS = 512; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_EXPERTS = 1024; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MIN_TOP_EXPERTS = 9; +static constexpr int kHIGH_EXPERT_LANE_OWNED_TOPK_MAX_TOP_EXPERTS = 16; + +template +using HighExpertLaneOwnedTopKRange = + LaneOwnedTopKRange; + +template +__forceinline__ __device__ void reduceTopKForLane( + cg::thread_block_tile const& warp, Type& out, int32_t& outIdx, + Type (&value)[N], int32_t (&idx)[N], Type const minValue, int32_t laneIdx) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K <= kWARP_SIZE, "Top K must have K <= kWARP_SIZE"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert(N <= 64, + "Only support candidates number less than or equal to " + "64*32=2048"); + using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + Sort::run(topK); + + typename RedType::TypeCmp packedMax{}; + typename RedType::TypeCmp lanePacked{}; +#pragma unroll + for (int kk = 0; kk < K; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compVal; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; + } + packedMax = topK[0].reduce(warp); + if (laneIdx == kk) { + lanePacked = packedMax; + } + } + + if (laneIdx < K) { + RedType::unpack(out, outIdx, lanePacked); + } else { + out = minValue; + outIdx = -1; + } +} } // namespace reduce_topk } // namespace moe diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index 86ee5397fbc..78aed496bba 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -1086,4 +1086,4 @@ void moe_lora_align_block_size( has_expert_map); } }); -} \ No newline at end of file +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 5a9c91d4563..e79f22f1d73 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -276,6 +276,61 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size); +void fused_kimi_k3_mla_key_concat_kv_cache_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_key_concat_ds_mla_insert( + torch::stable::Tensor& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, + torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8, + torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& k_scale_inv, + torch::stable::Tensor const& v_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& q_scale_inv, + torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + +void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe, + torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe, + torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, int64_t cache_block_size, + std::optional position_ids, + std::optional cos_sin_cache); + void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( torch::stable::Tensor const& q, torch::stable::Tensor const& kv, torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache, @@ -315,6 +370,19 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional index_q_out, const std::string& kv_cache_dtype, bool skip_index_branch); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE +void fused_kda_decode( + torch::stable::Tensor const& x, torch::stable::Tensor const& weight, + std::optional bias, + torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g, + torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, torch::stable::Tensor& state, + torch::stable::Tensor& out, std::optional lower_bound, + std::optional output_gate, + std::optional norm_weight, double norm_eps); +#endif + #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES void kimi_k3_attn_res(torch::stable::Tensor& prefix, torch::stable::Tensor const& delta, @@ -383,6 +451,20 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, torch::stable::Tensor& out, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); +void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t local_buffer, + fptr_t multicast_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); +void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, + fptr_t local_buffer, fptr_t epoch_buffer, + int64_t stage_sz_bytes); void dispose(fptr_t _fa); int64_t meta_size(); void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); @@ -421,6 +503,12 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double threshold); void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, double alpha = 1.702, double limit = 7.0); +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input, + double beta = 1.0, double linear_beta = -1.0); +void masked_situ_and_mul(torch::stable::Tensor& out, + torch::stable::Tensor& input, + const torch::stable::Tensor& expert_num_tokens, + double beta = 1.0, double linear_beta = -1.0); void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input); void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input); @@ -497,6 +585,13 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, const std::string& kv_cache_dtype, torch::stable::Tensor& scale); +void concat_and_cache_mla_grouped(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache_ptrs, + torch::stable::Tensor& slot_mapping, + int64_t block_size, int64_t block_stride, + int64_t entry_stride); + // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla void concat_and_cache_mla_rope_fused( torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index c364e211474..fec337a8953 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -324,7 +324,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). // conditionally compiled so impl registration is in source file ops.def( - "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " + "bool enable_pdl=False) -> ()"); // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). // conditionally compiled so impl registration is in source file @@ -447,6 +448,48 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, " "int cache_block_size) -> ()"); + // Kimi-K3 MLA epilogues: optional RoPE followed by concat/cache insertion. + ops.def( + "fused_kimi_k3_mla_key_concat_kv_cache_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_key_concat_ds_mla_insert(" + "Tensor! q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, " + "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" + "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " + "Tensor! q_fp8, Tensor! k_fp8, Tensor! v_fp8, Tensor! k_cache, " + "Tensor slot_mapping, Tensor q_scale_inv, Tensor k_scale_inv, " + "Tensor v_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + + // Kimi-K3 MLA decode epilogue: concat mqa_q = [ql_nope | q_pe] and insert the + // latent [kv_c_normed | k_pe] into the paged cache (bf16 / fp8 / fp8_ds_mla). + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "Tensor q_scale_inv, Tensor cache_scale_inv, int cache_block_size, " + "Tensor? position_ids=None, Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(" + "Tensor ql_nope, Tensor q_pe, Tensor kv_c_normed, Tensor k_pe, " + "Tensor! mqa_q, Tensor! k_cache, Tensor slot_mapping, " + "int cache_block_size, Tensor? position_ids=None, " + "Tensor? cos_sin_cache=None) -> ()"); + #ifndef USE_ROCM ops.def( "minimax_allreduce_rms_qk(" @@ -468,6 +511,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int block_size, Tensor!? q_out, Tensor!? index_q_out, " "str kv_cache_dtype, bool skip_index_branch=False) -> ()"); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.def( + "fused_kda_decode(" + "Tensor x, Tensor weight, Tensor? bias, Tensor! conv_state, " + "Tensor raw_g, Tensor raw_beta, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor! state, Tensor! out, " + "float? lower_bound=None, Tensor? output_gate=None, " + "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); +#endif + #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES ops.def( "kimi_k3_attn_res(" @@ -538,6 +591,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "limit=7.0) " "-> ()"); + // SituGLU implementation used in Kimi models. + ops.def( + "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float " + "linear_beta=-1.0) -> ()"); + ops.def( + "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " + "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); + // GELU implementation used in GPT-2. ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); @@ -696,11 +757,27 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl( "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert)); + ops.impl("fused_kimi_k3_mla_decode_q_concat_ds_mla_insert", + TORCH_BOX(&fused_kimi_k3_mla_decode_q_concat_ds_mla_insert)); #ifndef USE_ROCM ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); #endif ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); +#ifdef VLLM_ENABLE_FUSED_KDA_DECODE + ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); +#endif + #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); #endif @@ -726,6 +803,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("gelu_tanh_and_mul", TORCH_BOX(&gelu_tanh_and_mul)); ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul)); ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); + ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); + ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); ops.impl("gelu_quick", TORCH_BOX(&gelu_quick)); @@ -823,6 +902,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " str kv_cache_dtype," " Tensor scale) -> ()"); + // Grouped concat_and_cache_mla across all layers (bf16 only). Each + // layer's cache base pointer is read from kv_cache_ptrs. + ops.def( + "concat_and_cache_mla_grouped(Tensor kv_c, Tensor k_pe," + " Tensor kv_cache_ptrs," + " Tensor slot_mapping," + " int block_size, int block_stride," + " int entry_stride) -> ()"); + // Rotate Q and K, then write to kv cache for MLA ops.def( "concat_and_cache_mla_rope_fused(" @@ -921,6 +1009,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_grouped", + TORCH_BOX(&concat_and_cache_mla_grouped)); ops.impl("concat_and_cache_mla_rope_fused", TORCH_BOX(&concat_and_cache_mla_rope_fused)); ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 6e95d72ed5e..296978f384d 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -13,10 +13,18 @@ namespace vllm { namespace fp8 { #ifdef ENABLE_FP8 +// Unspecialized conversions are a compile error: the old passthrough +// (`return x;`) silently skipped fp8 encoding for any (Tout, Tin) pair +// without a specialization below (e.g. the torch stable-ABI scalar types), +// corrupting quantized data with no runtime signal. +template +inline constexpr bool _no_conversion_specialization = false; + template __inline__ __device__ Tout vec_conversion( const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) { - return x; + static_assert(_no_conversion_specialization, + "no vec_conversion specialization for this (Tout, Tin) pair"); } // float -> c10::Float8_e4m3fn @@ -301,7 +309,9 @@ __inline__ __device__ bf16_8_t vec_conversion( template __inline__ __device__ Tout scaled_vec_conversion( const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) { - return x; + static_assert( + _no_conversion_specialization, + "no scaled_vec_conversion specialization for this (Tout, Tin) pair"); } // fp8 -> half @@ -492,6 +502,25 @@ __inline__ __device__ uint8_t scaled_vec_conversion( __builtin_unreachable(); // Suppress missing return statement warning } +// torch stable-ABI (headeronly) scalar types delegate to the CUDA-native +// conversions, so libtorch_stable kernels dispatched on c10::BFloat16 / +// c10::Half quantize correctly without manual casts. +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::BFloat16& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const c10::Half& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + return scaled_vec_conversion( + reinterpret_cast(a), scale, fp8_type); +} + // float -> fp8 template <> __inline__ __device__ uint8_t scaled_vec_conversion( diff --git a/setup.py b/setup.py index 40d4ca103be..baa359d0b9d 100644 --- a/setup.py +++ b/setup.py @@ -783,6 +783,7 @@ class precompiled_wheel_utils: "vllm/_qutlass_C.abi3.so", "vllm/_flashmla_C.abi3.so", "vllm/_flashmla_extension_C.abi3.so", + "vllm/_flashkda_C.abi3.so", "vllm/_sparse_flashmla_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so", "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", @@ -1150,6 +1151,10 @@ if _is_cuda(): ext_modules.append( CMakeExtension(name="vllm._flashmla_extension_C", optional=True) ) + if USE_PRECOMPILED_EXTENSIONS or ( + CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0") + ): + ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True)) if envs.VLLM_USE_PRECOMPILED or ( CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") ): diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index edddb6ec845..5d29ad692fc 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random +from types import SimpleNamespace +from unittest.mock import Mock import pytest import ray @@ -9,6 +11,8 @@ import torch import torch.distributed as dist from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from vllm.distributed.device_communicators import custom_all_reduce +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce from vllm.distributed.parallel_state import get_tp_group, graph_capture from ..utils import ( @@ -23,6 +27,59 @@ for i, v in enumerate(test_sizes): test_sizes[i] -= v % 8 +def test_sp16_dispatches_only_to_mnnvl_lamport( + monkeypatch: pytest.MonkeyPatch, +): + """SP16 uses the MNNVL Lamport kernels and rejects same-host dispatch.""" + comm = CustomAllreduce.__new__(CustomAllreduce) + comm.disabled = False + comm.world_size = 16 + comm.fully_connected = False + comm.mnnvl_only = True + comm._IS_CAPTURING = False + comm.max_mnnvl_all_gather_size = 2 * 1024 * 1024 + comm.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024 + comm.mnnvl_multicast_ptr = 1 + comm.mnnvl_lamport_ag_local_ptr = 1 + comm.mnnvl_lamport_ag_multicast_ptr = 1 + comm.mnnvl_lamport_ag_epoch_ptr = 1 + comm.mnnvl_lamport_rs_local_ptr = 1 + comm.mnnvl_lamport_rs_epoch_ptr = 1 + comm.mnnvl_buffer_size = 32 * 1024 * 1024 + comm._ptr = 0 + + lamport_all_gather = Mock() + lamport_reduce_scatter = Mock() + monkeypatch.setattr(custom_all_reduce.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + custom_all_reduce, + "ops", + SimpleNamespace( + mnnvl_lamport_all_gather=lamport_all_gather, + mnnvl_lamport_reduce_scatter=lamport_reduce_scatter, + ), + ) + + gathered = comm.custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) + scattered = comm.custom_reduce_scatter(torch.empty((16, 8), dtype=torch.bfloat16)) + + assert gathered is not None + assert scattered is not None + lamport_all_gather.assert_called_once() + lamport_reduce_scatter.assert_called_once() + assert not comm.should_custom_ar(torch.empty(8, dtype=torch.bfloat16)) + assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.int32)) + assert not comm.should_custom_all_gather( + torch.empty((131073, 8), dtype=torch.bfloat16) + ) + + comm.mnnvl_only = False + assert not comm.should_custom_all_gather(torch.empty((8, 8), dtype=torch.bfloat16)) + assert not comm.should_custom_reduce_scatter( + torch.empty((16, 8), dtype=torch.bfloat16) + ) + + @ray.remote(num_gpus=1, max_calls=1) def graph_allreduce( monkeypatch: pytest.MonkeyPatch, @@ -80,6 +137,32 @@ def graph_allreduce( torch.testing.assert_close(out1, inp1) torch.testing.assert_close(out2, inp2) + fa = get_tp_group().device_communicator.ca_comm + tp_rank = rank % tp_size + with graph_capture(device=device) as graph_capture_context: + local = torch.full( + (512, 4096), tp_rank + 1, dtype=torch.bfloat16, device=device + ) + reduce_input = torch.full( + (512 * tp_size, 4096), + tp_rank + 1, + dtype=torch.bfloat16, + device=device, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + gathered = fa.custom_all_gather(local) + scattered = fa.custom_reduce_scatter(reduce_input) + graph.replay() + assert gathered is not None + assert scattered is not None + expected_gather = torch.cat( + [torch.full_like(local, peer_rank + 1) for peer_rank in range(tp_size)] + ) + expected_scatter = torch.full_like(local, tp_size * (tp_size + 1) // 2) + torch.testing.assert_close(gathered, expected_gather) + torch.testing.assert_close(scattered, expected_scatter) + @ray.remote(num_gpus=1, max_calls=1) def eager_allreduce( @@ -110,6 +193,29 @@ def eager_allreduce( out = fa.all_reduce(out, registered=False) torch.testing.assert_close(out, inp * (tp_size**num_communication)) + group = get_tp_group().device_group + tp_rank = rank % tp_size + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + local = torch.full((64, 4096), tp_rank + 1, dtype=dtype, device=device) + expected_gather = torch.empty( + (64 * tp_size, 4096), dtype=dtype, device=device + ) + dist.all_gather_into_tensor(expected_gather, local, group=group) + gathered = fa.custom_all_gather(local) + assert gathered is not None + torch.testing.assert_close(gathered, expected_gather) + + reduce_input = torch.full( + (64 * tp_size, 4096), tp_rank + 1, dtype=dtype, device=device + ) + expected_scatter = torch.empty((64, 4096), dtype=dtype, device=device) + dist.reduce_scatter_tensor( + expected_scatter, reduce_input.clone(), group=group + ) + scattered = fa.custom_reduce_scatter(reduce_input) + assert scattered is not None + torch.testing.assert_close(scattered, expected_scatter) + inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device) out = inp for _ in range(num_communication): @@ -130,3 +236,14 @@ def test_custom_allreduce( if world_size > torch.accelerator.device_count(): pytest.skip("Not enough GPUs to run the test.") multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) + + +@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) +def test_custom_collectives_world_size_four( + monkeypatch: pytest.MonkeyPatch, + test_target, +): + """Exercise the four-rank kernel specialization used by Kimi SP.""" + if torch.accelerator.device_count() < 4: + pytest.skip("Not enough GPUs to run the test.") + multi_process_parallel(monkeypatch, 4, 1, test_target) diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py new file mode 100644 index 00000000000..0e66e0da1e0 --- /dev/null +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RoPE equivalence tests for the fused Kimi-K3 MLA epilogues.""" + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA" +) + +_DTYPE = torch.bfloat16 +_NUM_TOKENS = 3 +_NUM_HEADS = 4 +_BLOCK_SIZE = 8 +_POSITIONS = (1, 7, 13) +_SLOTS = (0, 3, 9) + + +def _randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2 + + +def _rope_cache(max_position: int = 32) -> torch.Tensor: + inv_freq = 1.0 / ( + 50000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64) + ) + positions = torch.arange(max_position, dtype=torch.float32, device="cuda") + freqs = torch.outer(positions, inv_freq) + # The fused epilogue reads the cos/sin table in fp32 (RoPE math runs in fp32). + return torch.cat((freqs.cos(), freqs.sin()), dim=-1) + + +def _apply_gptj_rope( + x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor +) -> torch.Tensor: + cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1) + for _ in range(x.ndim - 2): + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + x1 = x[..., ::2].float() + x2 = x[..., 1::2].float() + out1 = x1 * cos.float() - x2 * sin.float() + out2 = x2 * cos.float() + x1 * sin.float() + return torch.stack((out1, out2), dim=-1).flatten(-2).to(x.dtype) + + +def _cache_rows(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor: + return cache.reshape(-1, cache.shape[-1]).index_select(0, slots) + + +def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close( + actual.float(), + expected.to(torch.float8_e4m3fn).float(), + atol=0.03125, + rtol=0.15, + ) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(0) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + q = _randn(_NUM_TOKENS, _NUM_HEADS, 192) + k_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + k_pe = _randn(_NUM_TOKENS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + v = _randn(_NUM_TOKENS, _NUM_HEADS, 128) + + q_expected = q.clone() + q_expected[..., 128:] = _apply_gptj_rope( + q_expected[..., 128:], positions, cos_sin_cache + ) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + k_expected = torch.cat( + (k_nope, k_pe_expected[:, None, :].expand(-1, _NUM_HEADS, -1)), dim=-1 + ) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = q.clone() + k_actual = fused_mla_key_concat_kv_cache_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros( + 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn + ) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual, k_actual, v_actual = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c, + v, + cache, + slots, + one, + one, + one, + one, + positions, + cos_sin_cache, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(k_actual, k_expected) + _assert_fp8_close(v_actual, v) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = q.clone() + k_actual = fused_mla_key_concat_ds_mla_insert( + q_actual, + k_nope, + k_pe, + kv_c, + cache, + slots, + positions, + cos_sin_cache, + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(k_actual, k_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) +@torch.inference_mode() +def test_decode_epilogue_fuses_gptj_rope(cache_kind: str) -> None: + torch.manual_seed(1) + positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + cos_sin_cache = _rope_cache() + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + + q_pe_expected = _apply_gptj_rope(q_pe, positions, cos_sin_cache) + k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache) + q_expected = torch.cat((ql_nope, q_pe_expected), dim=-1) + cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1) + + kwargs = {"positions": positions, "cos_sin_cache": cos_sin_cache} + if cache_kind == "bf16": + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots, **kwargs + ) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(_cache_rows(cache, slots), cache_expected) + elif cache_kind == "fp8": + cache = torch.zeros( + 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn + ) + one = torch.ones(1, device="cuda", dtype=torch.float32) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c, + k_pe, + cache, + slots, + q_scale_inv=one, + cache_scale_inv=one, + **kwargs, + ) + _assert_fp8_close(q_actual, q_expected) + _assert_fp8_close(_cache_rows(cache, slots), cache_expected) + else: + cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8) + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots, ds_mla=True, **kwargs + ) + rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE) + torch.testing.assert_close(q_actual, q_expected) + torch.testing.assert_close(rope_cache, k_pe_expected) + + +@torch.inference_mode() +def test_decode_epilogue_preserves_nope_path() -> None: + torch.manual_seed(2) + slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64) + ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512) + q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64) + kv_c = _randn(_NUM_TOKENS, 512) + k_pe = _randn(_NUM_TOKENS, 64) + cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots + ) + + torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1)) + torch.testing.assert_close( + _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1) + ) diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index f698c385fee..f6fc0ab2670 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -197,6 +197,46 @@ def test_silu_and_mul_with_clamp( opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) +@pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) +@torch.inference_mode() +def test_masked_situ_and_mul( + default_vllm_config, + linear_beta: float, + dtype: torch.dtype, +) -> None: + """Masked SITU computes valid expert rows and preserves padded zeros.""" + device = CUDA_DEVICES[0] + num_experts, max_num_tokens, d = 4, 7, 512 + beta = 1.5 + input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device) + expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device) + output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device) + + torch.ops._C.masked_situ_and_mul( + output, input, expert_num_tokens, beta, linear_beta + ) + + gate, up = input.float().chunk(2, dim=-1) + expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if linear_beta > 0: + up = linear_beta * torch.tanh(up / linear_beta) + expected = (expected * up).to(dtype) + for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()): + torch.testing.assert_close( + output[expert, :num_tokens], + expected[expert, :num_tokens], + atol=get_default_atol(output), + rtol=get_default_rtol(output), + ) + assert torch.count_nonzero(output[expert, num_tokens:]) == 0 + + opcheck( + torch.ops._C.masked_situ_and_mul, + (output, input, expert_num_tokens, beta, linear_beta), + ) + + @pytest.mark.parametrize( "activation", [ diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index b6a70b19b03..e5f2c878ff2 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -13,7 +13,7 @@ from __future__ import annotations import pytest import torch -from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm +from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 69788e37721..cccec939e3b 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -7,7 +7,9 @@ matching the eager triton kernel output.""" import pytest import torch -from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import ( + FusedRMSNormGated, +) from vllm.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16] diff --git a/tests/kernels/moe/test_grouped_topk.py b/tests/kernels/moe/test_grouped_topk.py index c58c8474b06..1ab8550d64e 100644 --- a/tests/kernels/moe/test_grouped_topk.py +++ b/tests/kernels/moe/test_grouped_topk.py @@ -23,6 +23,53 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +def _run_single_group_topk( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + return fused_grouped_topk( + hidden_states=torch.empty( + (logits.shape[0], 0), dtype=logits.dtype, device=logits.device + ), + gating_output=logits, + topk=topk, + renormalize=renormalize, + e_score_correction_bias=bias, + num_expert_group=1, + topk_group=1, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ) + + +def _single_group_reference( + logits: torch.Tensor, + bias: torch.Tensor, + topk: int, + *, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + if scoring_func == "sigmoid": + scores = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 + else: + scores = torch.softmax(logits, dim=-1).float() + indices = torch.argsort( + scores + bias.float(), dim=-1, descending=True, stable=True + )[:, :topk] + values = scores.gather(1, indices) + if renormalize: + values /= values.sum(dim=-1, keepdim=True) + 1e-20 + values *= routed_scaling_factor + return values, indices.to(torch.int32) + + @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) @@ -101,3 +148,222 @@ def test_grouped_topk( baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0 ) torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_grouped_topk_single_group_large_batch(): + set_random_seed(0) + logits = torch.randn((1536, 896), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((896,), dtype=torch.float32, device="cuda") + + expected_values, expected_ids = _single_group_reference( + logits, bias, 16, scoring_func="sigmoid", renormalize=True + ) + actual_values, actual_ids = _run_single_group_topk( + logits, bias, 16, scoring_func="sigmoid", renormalize=True + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize( + "num_experts,topk,input_dtype,bias_dtype", + [ + (512, 9, torch.bfloat16, torch.float32), + (512, 16, torch.float16, torch.float16), + (513, 9, torch.float32, torch.bfloat16), + (513, 16, torch.bfloat16, torch.float32), + (895, 9, torch.float16, torch.bfloat16), + (896, 16, torch.float32, torch.float16), + (897, 9, torch.bfloat16, torch.bfloat16), + (897, 16, torch.float16, torch.float32), + (1024, 9, torch.float32, torch.bfloat16), + (1024, 16, torch.bfloat16, torch.float16), + ], +) +@pytest.mark.parametrize( + "scoring_func,renormalize,routed_scaling_factor", + [ + ("sigmoid", True, 1.0), + ("sigmoid", False, 2.5), + ("softmax", True, 2.5), + ("softmax", False, 1.0), + ], +) +def test_grouped_topk_single_group_tiers( + num_experts: int, + topk: int, + input_dtype: torch.dtype, + bias_dtype: torch.dtype, + scoring_func: str, + renormalize: bool, + routed_scaling_factor: float, +): + set_random_seed(7) + logits = torch.randn((17, num_experts), dtype=input_dtype, device="cuda") + bias = torch.randn((num_experts,), dtype=bias_dtype, device="cuda") + + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize( + "num_experts,topk,scoring_func", + [ + (128, 8, "sigmoid"), + (129, 8, "sigmoid"), + (257, 8, "sigmoid"), + (385, 8, "sigmoid"), + (512, 9, "sigmoid"), + (513, 9, "sigmoid"), + (769, 9, "sigmoid"), + (897, 16, "sigmoid"), + (1024, 16, "sigmoid"), + (128, 4, "softmax"), + (128, 5, "softmax"), + (129, 8, "softmax"), + (161, 8, "softmax"), + (256, 9, "softmax"), + (257, 8, "softmax"), + (512, 9, "softmax"), + (512, 17, "softmax"), + (512, 23, "softmax"), + (513, 8, "softmax"), + (577, 9, "softmax"), + (769, 9, "softmax"), + (897, 9, "softmax"), + (1024, 16, "softmax"), + ], +) +def test_grouped_topk_single_group_capacity_tiers( + num_experts: int, + topk: int, + scoring_func: str, +): + set_random_seed(11) + logits = torch.randn((3, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + expected_values, expected_ids = _single_group_reference( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + topk, + scoring_func=scoring_func, + renormalize=True, + routed_scaling_factor=2.5, + ) + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +def test_grouped_topk_single_group_stable_ties(num_experts: int): + logits = torch.zeros((1, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=True, + routed_scaling_factor=2.5, + ) + + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + expected_values = torch.full((1, 16), 2.5 / 16, dtype=torch.float32, device="cuda") + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +@pytest.mark.parametrize("num_experts", [512, 896, 1024]) +@pytest.mark.parametrize("num_finite", [0, 15]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_grouped_topk_single_group_nonfinite_scores( + num_experts: int, num_finite: int, renormalize: bool +): + logits = torch.full( + (1, num_experts), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + if num_finite: + logits[0, :num_finite] = torch.arange( + num_finite, dtype=torch.bfloat16, device="cuda" + ) + logits[0, num_finite] = torch.inf + logits[0, num_finite + 1] = -torch.inf + bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") + + actual_values, actual_ids = _run_single_group_topk( + logits, + bias, + 16, + scoring_func="sigmoid", + renormalize=renormalize, + routed_scaling_factor=2.5, + ) + + if num_finite == 0: + expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] + if renormalize: + expected_values = torch.full( + (1, 16), 1 / 16, dtype=torch.float32, device="cuda" + ) + else: + expected_values = torch.zeros((1, 16), dtype=torch.float32, device="cuda") + else: + expected_ids = torch.cat( + ( + torch.arange(num_finite - 1, -1, -1, dtype=torch.int32, device="cuda"), + torch.tensor([num_finite], dtype=torch.int32, device="cuda"), + ) + )[None] + finite_values = logits[0, :num_finite].float().sigmoid().flip(0) + if renormalize: + finite_values /= finite_values.sum() + finite_values *= 2.5 + expected_values = torch.cat( + (finite_values, torch.zeros(1, dtype=torch.float32, device="cuda")) + )[None] + + torch.testing.assert_close(actual_ids, expected_ids) + torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index a017fa07bf9..7d07af4080c 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -269,6 +269,7 @@ def test_moe_align_block_size_with_expert_map( if (experts[k] in local_experts) or not mask_inactive_experts else -1 ) + topk_ids[0, 0] = -1 actual_sorted_ids, actual_expert_ids, actual_num_tokens = moe_align_block_size( topk_ids=topk_ids, diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py new file mode 100644 index 00000000000..0d5813dd50a --- /dev/null +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -0,0 +1,649 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Kimi-K3 SM103 decode GEMM selector (shape-only dispatch).""" + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import regex as re +import torch +from torch import nn + +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import SkinnyGemmConfig +from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm +from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS + +# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is +# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16). +EXPECTED_SELECTIONS = { + (1536, 128): (set(), set(range(1, 17))), + (3072, 128): (set(), set(range(1, 17))), + (1536, 7168): (set(), set(range(1, 17))), + (3072, 7168): (set(range(1, 6)), set()), + (2112, 7168): (set(), set(range(1, 17))), + (2304, 1536): (set(), set(range(1, 17))), + (4608, 1536): (set(), set(range(1, 17))), + (3584, 7168): ({1}, set(range(2, 9))), + (6288, 7168): (set(range(1, 5)), set()), + (12448, 7168): (set(range(1, 4)), set()), + (7168, 768): (set(), set(range(1, 17))), + (7168, 1536): ({1}, set()), + (7168, 3072): ({1, 2}, set()), + (7168, 3584): ({1, 2}, set()), + (7168, 4224): ({1}, set()), + (7168, 8448): (set(range(1, 4)), set()), + (8448, 7168): ({1, 2}, set()), + (16896, 7168): ({1, 2}, set()), + (20480, 7168): (set(range(1, 5)), set()), + (40960, 7168): (set(range(1, 5)), set()), + # TP16. + (3216, 7168): (set(range(1, 6)), set(range(9, 16))), + (768, 7168): (set(range(1, 5)), set(range(5, 17))), + (1152, 1536): ({1}, set(range(2, 17))), + (768, 128): (set(), set(range(1, 17))), + (7168, 384): (set(), set(range(1, 9))), + (4224, 7168): (set(range(1, 4)), set(range(4, 9))), + (10240, 7168): (set(range(1, 5)), set()), +} + +CUTE_CASES = [ + (spec.n, spec.k, num_tokens) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, _ in spec.cute_configs +] + +RESIDUAL_CUTE_CASES = [ + (spec.n, spec.k, num_tokens) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, _ in spec.residual_configs +] + +EXPECTED_CUTE_CONFIGS = { + (3072, 7168, 1): (224, 3, 4, 8), + (3072, 7168, 2): (128, 3, 2, 8), + (3072, 7168, 3): (128, 2, 1, 8), + (3072, 7168, 4): (64, 2, 2, 8), + (3072, 7168, 5): (128, 3, 1, 8), + (3584, 7168, 1): (224, 2, 4, 8), + (6288, 7168, 1): (224, 3, 4, 8), + (6288, 7168, 2): (64, 3, 2, 8), + (6288, 7168, 3): (32, 3, 4, 8), + (6288, 7168, 4): (128, 6, 1, 8), + (12448, 7168, 1): (224, 4, 2, 8), + (12448, 7168, 2): (64, 4, 2, 8), + (12448, 7168, 3): (64, 2, 2, 8), + (7168, 1536, 1): (96, 4, 2, 8), + (7168, 3072, 1): (96, 2, 4, 8), + (7168, 3072, 2): (32, 4, 4, 8), + (7168, 3584, 1): (224, 4, 2, 8), + (7168, 3584, 2): (64, 4, 2, 8), + (7168, 4224, 1): (96, 4, 2, 4), + (7168, 8448, 1): (32, 4, 4, 8), + (7168, 8448, 2): (96, 4, 1, 8), + (7168, 8448, 3): (96, 4, 1, 8), + (8448, 7168, 1): (224, 3, 4, 8), + (8448, 7168, 2): (32, 4, 4, 8), + (16896, 7168, 1): (224, 6, 4, 8), + (16896, 7168, 2): (32, 4, 4, 8), + (20480, 7168, 1): (224, 4, 2, 8), + (20480, 7168, 2): (64, 4, 2, 8), + (20480, 7168, 3): (64, 2, 2, 8), + (20480, 7168, 4): (64, 4, 1, 8), + (40960, 7168, 1): (128, 4, 2, 8), + (40960, 7168, 2): (64, 4, 2, 8), + (40960, 7168, 3): (64, 2, 2, 8), + (40960, 7168, 4): (64, 4, 1, 8), + # TP16. + (3216, 7168, 1): (224, 3, 4, 8), + (3216, 7168, 2): (128, 4, 2, 8), + (3216, 7168, 3): (128, 2, 1, 8), + (3216, 7168, 4): (64, 2, 2, 8), + (3216, 7168, 5): (128, 3, 1, 8), + (768, 7168, 1): (224, 2, 4, 8), + (768, 7168, 2): (224, 2, 2, 8), + (768, 7168, 3): (224, 2, 2, 8), + (768, 7168, 4): (224, 2, 2, 8), + (1152, 1536, 1): (192, 3, 4, 8), + (4224, 7168, 1): (224, 3, 4, 8), + (4224, 7168, 2): (128, 2, 1, 8), + (4224, 7168, 3): (64, 2, 2, 8), + (10240, 7168, 1): (224, 4, 2, 8), + (10240, 7168, 2): (32, 2, 4, 8), + (10240, 7168, 3): (64, 4, 1, 8), + (10240, 7168, 4): (64, 4, 1, 8), +} + +EXPECTED_RESIDUAL_CUTE_CONFIGS = { + (7168, 3584, 1): (64, 4, 2, 8), + (7168, 3584, 2): (64, 7, 2, 8), + (7168, 3584, 3): (64, 2, 1, 8), + (7168, 3584, 4): (64, 2, 1, 8), +} + + +def _config_tuple(config) -> tuple[int, int, int, int]: + return ( + config.block_size, + config.outputs_per_block, + config.k_unroll, + config.vector_width, + ) + + +def test_table_is_keyed_by_shape() -> None: + for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items(): + assert (spec.n, spec.k) == (n, k) + + +def test_every_dsv3_routed_shape_is_instantiated() -> None: + """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises. + + The table routes by shape while the kernel is built per shape, so a missing + instantiation only shows up at the token counts that route to dsv3. Checking + it here needs no GPU, which is the point -- a GPU-only check is exactly what + let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). + """ + source = ( + Path(__file__).resolve().parents[2] + / "csrc" + / "libtorch_stable" + / "dsv3_fused_a_gemm.cu" + ).read_text(encoding="utf-8") + # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built + # by default, so they must not count as available. + production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0] + explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split( + "#ifdef VLLM_K3_BENCH_SHAPES" + )[0] + compiled = { + (int(hd_in), int(hd_out)) + for hd_in, hd_out in re.findall( + r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros + ) + } | { + (int(hd_in), int(hd_out)) + for hd_in, hd_out in re.findall(r"hd_in == (\d+) && hd_out == (\d+)", explicit) + } + assert compiled, "failed to parse the dispatch list" + + missing = sorted( + (spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + if spec.dsv3_tokens and (spec.k, spec.n) not in compiled + ) + assert not missing, ( + f"routed to dsv3 with no instantiation: {missing}; add " + "DISPATCH_DSV3_SHAPE(K, N) for each" + ) + + +def test_packed_row_major_rejects_single_row_slice() -> None: + packed = torch.empty(1, 128) + sliced = torch.empty(1, 144)[:, :128] + + assert packed.is_contiguous() + assert sliced.is_contiguous() + assert k3_gemm._is_packed_row_major(packed) + assert not k3_gemm._is_packed_row_major(sliced) + + +def test_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.cute_configs + } + assert actual == EXPECTED_CUTE_CONFIGS + + +def test_residual_cute_configs_match_measured_table() -> None: + actual = { + (spec.n, spec.k, num_tokens): _config_tuple(config) + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() + for num_tokens, config in spec.residual_configs + } + assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_sm103_selector_table(key: tuple[int, int]) -> None: + n, k = key + cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key] + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k) + if num_tokens in cute_tokens: + assert backend == "cute" + elif num_tokens in dsv3_tokens: + assert backend == "dsv3_fused_a" + else: + assert backend is None + + +@pytest.mark.parametrize("key", EXPECTED_SELECTIONS) +def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None: + n, k = key + assert k3_gemm.select_kimi_k3_backend(0, n, k) is None + assert k3_gemm.select_kimi_k3_backend(17, n, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None + assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None + + +def test_unlisted_shape_and_unselected_tokens_fall_back() -> None: + # Shape absent from the table. + assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None + # o_proj (7168,1536) is CuTe M1 only; M2+ falls back. + assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_sm103_residual_selector_table(num_tokens: int) -> None: + backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True) + assert backend == ("cute" if num_tokens <= 4 else None) + + +def test_build_plan_matches_selector() -> None: + for spec in k3_gemm.KIMI_K3_PROJECTIONS.values(): + plan = k3_gemm._build_plan(spec) + for num_tokens in range(1, 17): + backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k) + if backend is None: + assert num_tokens not in plan + else: + assert plan[num_tokens][0] == backend + + +def test_installation_is_shape_specific_and_unquantized( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeLinear(nn.Module): + def __init__(self, quant_method: object, n: int, k: int) -> None: + super().__init__() + self.quant_method = quant_method + self.weight = torch.empty(n, k) + + class FakeHead(nn.Module): + def __init__(self, n: int, k: int) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedEmbeddingMethod() + self.weight = torch.empty(n, k) + + root = nn.Module() + # dsv3-only shape (no cute warmup contribution). + root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536) + # quantized: must be left untouched. + quantized_method = object() + root.quantized = FakeLinear(quantized_method, 6288, 7168) + # cute shape. + root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168) + # cute + residual shape. + root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584) + # shape absent from the table: must be left untouched. + root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678) + root.lm_head = FakeHead(20480, 7168) + + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True) + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True) + + def request_warmup_configs(dtype, configs, *, has_residual=False): + target = residual_warmup_configs if has_residual else warmup_configs + target.update(configs) + + monkeypatch.setattr( + k3_gemm.shape_dynamic_skinny_gemm, + "request_warmup_configs", + request_warmup_configs, + ) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16) + + assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) + assert root.quantized.quant_method is quantized_method + assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod + assert isinstance( + root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod + ) + # Warmup covers only the installed modules' local (N, K). + assert warmup_configs == { + config + for key in ((6288, 7168), (7168, 3584), (20480, 7168)) + for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs + } + assert residual_warmup_configs == { + config + for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs + } + + +@pytest.mark.parametrize( + "dtype,platform_enabled", + [(torch.float16, True), (torch.bfloat16, False)], +) +def test_installation_requires_bf16_sm103( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + platform_enabled: bool, +) -> None: + class FakeLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.quant_method = k3_gemm.UnquantizedLinearMethod() + self.weight = torch.empty(2304, 1536) + + root = nn.Module() + root.projection = FakeLinear() + monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) + monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled) + + k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype) + + assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod + + +def _require_sm103_and_dsv3() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + pytest.skip("dsv3_fused_a_gemm was not built") + + +def _require_sm103_and_cute() -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3): + pytest.skip("Kimi-K3 production selection requires SM103") + if not k3_gemm.shape_dynamic_skinny_gemm.is_available(): + pytest.skip("CuTe DSL is not available") + + +@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES) +def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]: + """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch.""" + if not tokens: + return set() + return {min(tokens), max(tokens)} | ({8, 9} & set(tokens)) + + +# Derived from the table rather than hand-listed, so a shape routed to dsv3 +# cannot be added without being exercised here. +DSV3_CASES = sorted( + (num_tokens, spec.n, spec.k) + for spec in KIMI_K3_PROJECTIONS.values() + for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens) +) + + +@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES) +def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None: + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + assert num_tokens in spec.dsv3_tokens + torch.manual_seed(42) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + output = k3_gemm.try_low_latency_gemm(x, weight) + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +def test_nonpacked_single_token_dsv3_falls_back() -> None: + _require_sm103_and_dsv3() + n, k = 1536, 128 + storage = torch.randn(1, k + 16, dtype=torch.bfloat16, device="cuda") + x = storage[:, :k] + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + method = k3_gemm.KimiK3LowLatencyLinearMethod( + k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) + ) + + assert x.is_contiguous() + assert x.stride() == (k + 16, 1) + assert not k3_gemm._runtime_ok(x, weight) # strict guard rejects the view + output = method.apply(SimpleNamespace(weight=weight), x) + + reference = torch.nn.functional.linear(x, weight) + torch.testing.assert_close(output, reference) + + +def test_selected_kernels_cuda_graph_capture() -> None: + _require_sm103_and_cute() + _require_sm103_and_dsv3() + cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)] + dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda") + cute_weight = torch.randn( + cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda" + ) + dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda") + dsv3_weight = torch.randn( + dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda" + ) + k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight) + dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) + graph.replay() + torch.accelerator.synchronize() + + assert cute_output is not None + assert dsv3_output is not None + for output, activation, weight in ( + (cute_output, cute_x, cute_weight), + (dsv3_output, dsv3_x, dsv3_weight), + ): + reference = torch.nn.functional.linear(activation, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", [1, 8, 9, 16]) +def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None: + """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch.""" + _require_sm103_and_dsv3() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + k3_gemm.try_low_latency_gemm(x, weight) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.try_low_latency_gemm(x, weight) + graph.replay() + torch.accelerator.synchronize() + + assert output is not None + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES) +def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None: + _require_sm103_and_cute() + torch.manual_seed(42 + num_tokens) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] + config = spec.residual_config(num_tokens) + assert config is not None + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +@pytest.mark.parametrize("num_tokens", range(1, 17)) +def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None: + _require_sm103_and_cute() + from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + ShapeDynamicSkinnyGemm, + ) + + n, k = 64, 512 + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") + config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k) + + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + + reference = x.float() @ weight.float().t() + residual.float() + torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) + + +@pytest.mark.parametrize("num_tokens", range(1, 5)) +def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None: + _require_sm103_and_cute() + spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)] + config = spec.residual_config(num_tokens) + assert config is not None + x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") + residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda") + k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) + graph.replay() + torch.accelerator.synchronize() + + reference = x.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 + + +class _SkinnyGemmSpy: + """Wraps the skinny-GEMM singleton to record whether CuTe was invoked.""" + + def __init__(self, real: Any) -> None: + self._real = real + self.calls: list[int] = [] + + def __call__(self, a, b, config=None, residual=None): + self.calls.append(a.shape[0]) + return self._real(a, b, config, residual) + + def is_available(self) -> bool: + return self._real.is_available() + + +@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) +def test_latent_moe_production_layout_residual( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, +) -> None: + """The real Latent-MoE residual is a non-packed slice of a cat buffer. + + The strict packed-row-major guard rejects such a slice at every token count + (a size-1 leading dim reads as contiguous but its stride is not packed), so + the CuTe residual epilogue never fires for this production layout and the + method falls back to addmm. Output is correct regardless of the path. + """ + _require_sm103_and_cute() + latent_dim, shared_dim = 3584, 7168 # routed_expert_up_proj K, N + torch.manual_seed(7 + num_tokens) + buf = torch.randn( + num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda" + ) + latent = buf[:, :latent_dim] # non-contiguous view (row stride = full width) + residual = buf[:, latent_dim:] # non-contiguous view + weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda") + + spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)] + method = k3_gemm.KimiK3LowLatencyLinearMethod( + k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) + ) + spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm) + monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy) + + layer = SimpleNamespace(weight=weight) + output = method.apply_with_residual(layer, latent, residual) + + reference = latent.float() @ weight.float().t() + residual.float() + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.flatten(), dim=0 + ).item() + assert cosine > 0.999 # correct regardless of the path taken + assert not spy.calls, ( + "non-packed buf-slice residual must fall back to addmm at every M" + ) + + +def test_residual_dispatch_falls_back_to_addmm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fallback = torch.randn(2, 3) + residual = torch.randn(2, 3) + x = torch.randn(2, 4) + weight = torch.randn(3, 4) + monkeypatch.setattr(torch, "addmm", lambda *args: fallback) + # CPU tensors fail the runtime check, forcing the addmm fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual) + + assert output is fallback + + +def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None: + fallback = torch.empty(2, 8) + monkeypatch.setattr( + k3_gemm.UnquantizedLinearMethod, + "apply", + lambda *args: fallback, + ) + # 1-D input fails the runtime check, forcing the base-method fallback. + method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) + + output = method.apply( + SimpleNamespace(weight=torch.empty(0)), + torch.empty(0), + ) + + assert output is fallback diff --git a/tests/kernels/test_kda.py b/tests/kernels/test_kda.py deleted file mode 100644 index 75c553fb864..00000000000 --- a/tests/kernels/test_kda.py +++ /dev/null @@ -1,226 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Precision tests for vllm's chunk_kda Triton operator. - -Compares chunk_kda against a naive recurrent reference (float32). -Uses torch.rand for q/k/v to match FLA's test pattern. -""" - -import pytest -import torch -import torch.nn.functional as F - -from vllm.third_party.flash_linear_attention.ops.kda import ( - chunk_kda, - chunk_kda_with_fused_gate, - fused_kda_gate, -) -from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd - -DEVICE = "cuda" - - -def naive_recurrent_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Naive recurrent KDA reference, ported from FLA's naive.py.""" - dtype = v.dtype - B, T, H, K = q.shape - V = v.shape[-1] - if scale is None: - scale = K**-0.5 - - q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) - q = q * scale - - S = k.new_zeros(B, H, K, V).to(q) - if initial_state is not None: - S += initial_state - o = torch.zeros_like(v) - for i in range(T): - q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] - S = S * g_i[..., None].exp() - S = S + torch.einsum( - "bhk,bhv->bhkv", - b_i[..., None] * k_i, - v_i - (k_i[..., None] * S).sum(-2), - ) - o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) - if not output_final_state: - S = None - return o.to(dtype), S - - -def assert_close( - name: str, - ref: torch.Tensor, - tri: torch.Tensor, - ratio: float, - err_atol: float = 1e-6, -): - """RMSE-based relative error comparison.""" - abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() - rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() - rmse_base = ref.detach().flatten().square().mean().sqrt().item() - rel_err = rmse_diff / (rmse_base + 1e-8) - print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") - if abs_err <= err_atol: - return - assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" - assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" - assert rel_err < ratio, ( - f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" - ) - - -@pytest.mark.parametrize( - ("H", "D", "cu_seqlens", "dtype"), - [ - pytest.param( - *test, - id="H{}-D{}-cu{}-{}".format(*test), - ) - for test in [ - (32, 128, [0, 64], torch.float16), - (32, 128, [0, 1024], torch.float16), - (32, 128, [0, 15], torch.float16), - (32, 128, [0, 256, 512, 768, 1024], torch.float16), - (32, 128, [0, 15, 100, 300, 1200], torch.float16), - (64, 128, [0, 256, 500, 1000], torch.float16), - (32, 128, [0, 8192], torch.float16), - (32, 128, [0, 256, 500, 1000], torch.bfloat16), - ] - ], -) -@torch.inference_mode() -def test_chunk_kda( - H: int, - D: int, - cu_seqlens: list[int], - dtype: torch.dtype, -): - T = cu_seqlens[-1] - torch.manual_seed(42) - B = 1 - cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) - N = len(cu_seqlens) - 1 - - q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) - g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( - dtype - ) - beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() - h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) - - # Naive reference with l2norm_fwd (same kernel as chunk_kda) - ref_outputs = [] - ref_states = [] - for i in range(N): - s, e = cu_seqlens[i], cu_seqlens[i + 1] - q_i = l2norm_fwd(q[:, s:e].contiguous()) - k_i = l2norm_fwd(k[:, s:e].contiguous()) - o_i, ht_i = naive_recurrent_kda( - q_i, - k_i, - v[:, s:e], - g[:, s:e], - beta[:, s:e], - initial_state=h0[i], - output_final_state=True, - ) - ref_outputs.append(o_i) - ref_states.append(ht_i) - ref_o = torch.cat(ref_outputs, dim=1) - ref_ht = torch.cat(ref_states, dim=0) - - # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) - tri_o, tri_ht = chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=g.clone(), - beta=beta.clone(), - initial_state=h0.transpose(-1, -2).contiguous().clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - - assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" - assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" - assert_close("o", ref_o, tri_o, 0.005) - assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) - - -@pytest.mark.parametrize( - ("cu_seqlens", "dtype"), - [ - ([0, 64], torch.float16), - ([0, 15, 100, 300], torch.bfloat16), - ], -) -@torch.inference_mode() -def test_chunk_kda_fused_gate_cumsum_matches_unfused( - cu_seqlens: list[int], - dtype: torch.dtype, -): - H, D = 8, 64 - T = cu_seqlens[-1] - N = len(cu_seqlens) - 1 - torch.manual_seed(123) - - cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) - q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) - beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid() - A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() - dt_bias = ( - torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 - ).contiguous() - h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) - initial_state = h0.transpose(-1, -2).contiguous() - - gate = fused_kda_gate( - raw_g.reshape(T, H * D), - A_log, - D, - g_bias=dt_bias, - ).unsqueeze(0) - old_o, old_ht = chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=gate, - beta=beta.clone(), - initial_state=initial_state.clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - new_o, new_ht = chunk_kda_with_fused_gate( - q=q.clone(), - k=k.clone(), - v=v.clone(), - raw_g=raw_g, - beta=beta.clone(), - A_log=A_log, - g_bias=dt_bias, - initial_state=initial_state.clone(), - output_final_state=True, - cu_seqlens=cu_seqlens_t, - use_qk_l2norm_in_kernel=True, - ) - - assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) - assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py index 69c9213647a..38a34bc923a 100644 --- a/tests/models/kimi_k3/test_attn_res.py +++ b/tests/models/kimi_k3/test_attn_res.py @@ -5,6 +5,7 @@ import pytest import torch import torch.nn.functional as F +from vllm.models.kimi_k3.common.mtp import fused_mtp_input from vllm.models.kimi_k3.nvidia.ops import attn_res from vllm.platforms import current_platform @@ -191,3 +192,35 @@ def test_attn_res_without_output_norm(): ) torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + + +@pytest.mark.parametrize("num_tokens", [0, 1, 17]) +def test_fused_mtp_input(num_tokens: int): + positions = torch.arange(num_tokens, device="cuda") + inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7) + previous_hidden_states = _randn_with_row_padding( + num_tokens, HIDDEN_SIZE, padding=11 + ) + enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + + masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + expected = torch.cat( + ( + F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS), + F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS), + ), + dim=-1, + ) + actual = fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + EPS, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE) + assert actual.is_contiguous() diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py new file mode 100644 index 00000000000..61a24a83e04 --- /dev/null +++ b/tests/models/kimi_k3/test_eagle3.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +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 ( + KimiK3ForConditionalGeneration, + KimiLinearModel, +) + + +def _make_kimi_linear_model() -> KimiLinearModel: + model = object.__new__(KimiLinearModel) + object.__setattr__(model, "aux_hidden_state_layers", (2,)) + object.__setattr__(model, "use_sequence_parallel", False) + return model + + +def test_kimi_k3_advertises_eagle3_support(): + assert supports_eagle3(KimiK3ForConditionalGeneration) + + +def test_kimi_k3_uses_shared_eagle3_layer_configuration(): + target = object.__new__(KimiK3ForConditionalGeneration) + torch.nn.Module.__init__(target) + model = _make_kimi_linear_model() + object.__setattr__(model, "layers", [None] * 93) + language_model = SimpleNamespace( + embed_input_ids=lambda _: None, + model=model, + ) + object.__setattr__(target, "language_model", language_model) + object.__setattr__(target, "_language_model_names", ["language_model"]) + + target.set_aux_hidden_state_layers((2, 46, 90)) + + assert model.aux_hidden_state_layers == (2, 46, 90) + assert target.get_eagle3_default_aux_hidden_state_layers() == ( + 2, + 46, + 90, + ) + + +def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + layer_residual = torch.tensor([[5.0, 6.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, None, layer_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", False) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + expected_layer_output = layer_hidden_states + layer_residual + torch.testing.assert_close(output, expected_layer_output) + torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) + torch.testing.assert_close(aux_hidden_states[1], expected_layer_output) + + +def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + final_hidden_states = torch.tensor([[9.0, 10.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (0, 1)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + final_attn_res = Mock(return_value=final_hidden_states) + monkeypatch.setattr(kimi_model, "attn_res", final_attn_res) + + output, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + torch.testing.assert_close(output, final_hidden_states) + 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 diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py new file mode 100644 index 00000000000..be916f6c0d0 --- /dev/null +++ b/tests/models/kimi_k3/test_kda.py @@ -0,0 +1,757 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Precision tests for vllm's chunk_kda Triton operator. + +Compares chunk_kda against a naive recurrent reference (float32). +Uses torch.rand for q/k/v to match FLA's test pattern. +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update +from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from vllm.models.kimi_k3.nvidia.kda import ( + is_flashkda_supported, + is_fused_kda_decode_supported, +) +from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda, + chunk_kda_with_fused_gate, + fused_kda_gate, + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd + +DEVICE = "cuda" + + +@torch.inference_mode() +def test_gather_initial_states_correctness(): + row_size = 8 * 128 * 128 + storage = torch.randn(5, row_size + 256, dtype=torch.float32, device=DEVICE) + state = storage[:, :row_size].view(5, 8, 128, 128) + assert not state.is_contiguous() + assert state[0].is_contiguous() + indices = torch.tensor([4, 1, 3], dtype=torch.int32, device=DEVICE) + has_initial_state = torch.tensor([True, False, True], device=DEVICE) + + expected = state[indices].clone() + expected[~has_initial_state] = 0 + + torch.testing.assert_close( + gather_initial_states(state, indices, has_initial_state), + expected, + ) + + +def naive_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Naive recurrent KDA reference, ported from FLA's naive.py.""" + dtype = v.dtype + B, T, H, K = q.shape + V = v.shape[-1] + if scale is None: + scale = K**-0.5 + + q, k, v, g, beta = (x.to(torch.float) for x in [q, k, v, g, beta]) + q = q * scale + + S = k.new_zeros(B, H, K, V).to(q) + if initial_state is not None: + S += initial_state + o = torch.zeros_like(v) + for i in range(T): + q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] + S = S * g_i[..., None].exp() + S = S + torch.einsum( + "bhk,bhv->bhkv", + b_i[..., None] * k_i, + v_i - (k_i[..., None] * S).sum(-2), + ) + o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S) + if not output_final_state: + S = None + return o.to(dtype), S + + +def assert_close( + name: str, + ref: torch.Tensor, + tri: torch.Tensor, + ratio: float, + err_atol: float = 1e-6, +): + """RMSE-based relative error comparison.""" + abs_err = (ref.detach() - tri.detach()).flatten().abs().max().item() + rmse_diff = (ref.detach() - tri.detach()).flatten().square().mean().sqrt().item() + rmse_base = ref.detach().flatten().square().mean().sqrt().item() + rel_err = rmse_diff / (rmse_base + 1e-8) + print(f"{name:>4} | abs={abs_err:.6f} | rmse={rel_err:.6f} | thr={ratio}") + if abs_err <= err_atol: + return + assert not torch.isnan(ref).any(), f"{name}: NaN detected in ref" + assert not torch.isnan(tri).any(), f"{name}: NaN detected in tri" + assert rel_err < ratio, ( + f"{name}: max abs err {abs_err:.6f}, rmse ratio {rel_err:.6f} >= {ratio}" + ) + + +@pytest.mark.parametrize( + ("H", "D", "cu_seqlens", "dtype"), + [ + pytest.param( + *test, + id="H{}-D{}-cu{}-{}".format(*test), + ) + for test in [ + (32, 128, [0, 64], torch.float16), + (32, 128, [0, 1024], torch.float16), + (32, 128, [0, 15], torch.float16), + (32, 128, [0, 256, 512, 768, 1024], torch.float16), + (32, 128, [0, 15, 100, 300, 1200], torch.float16), + (64, 128, [0, 256, 500, 1000], torch.float16), + (32, 128, [0, 8192], torch.float16), + (32, 128, [0, 256, 500, 1000], torch.bfloat16), + ] + ], +) +@torch.inference_mode() +def test_chunk_kda( + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, +): + T = cu_seqlens[-1] + torch.manual_seed(42) + B = 1 + cu_seqlens_t = torch.LongTensor(cu_seqlens).to(DEVICE) + N = len(cu_seqlens) - 1 + + q = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + k = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + v = torch.rand(B, T, H, D, dtype=dtype, device=DEVICE) + g = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float32, device=DEVICE)).to( + dtype + ) + beta = torch.rand(B, T, H, dtype=dtype, device=DEVICE).sigmoid() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + + # Naive reference with l2norm_fwd (same kernel as chunk_kda) + ref_outputs = [] + ref_states = [] + for i in range(N): + s, e = cu_seqlens[i], cu_seqlens[i + 1] + q_i = l2norm_fwd(q[:, s:e].contiguous()) + k_i = l2norm_fwd(k[:, s:e].contiguous()) + o_i, ht_i = naive_recurrent_kda( + q_i, + k_i, + v[:, s:e], + g[:, s:e], + beta[:, s:e], + initial_state=h0[i], + output_final_state=True, + ) + ref_outputs.append(o_i) + ref_states.append(ht_i) + ref_o = torch.cat(ref_outputs, dim=1) + ref_ht = torch.cat(ref_states, dim=0) + + # h0 transposed to (V, K) layout for the kernel; naive uses (K, V) + tri_o, tri_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + initial_state=h0.transpose(-1, -2).contiguous().clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert not torch.isnan(tri_o).any(), "Triton output o contains NaN" + assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN" + assert_close("o", ref_o, tri_o, 0.005) + assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005) + + +@pytest.mark.parametrize( + ("cu_seqlens", "dtype", "lower_bound"), + [ + ([0, 64], torch.float16, None), + ([0, 15, 100, 300], torch.bfloat16, None), + ([0, 15, 100, 300], torch.bfloat16, -3.0), + ], +) +@torch.inference_mode() +def test_chunk_kda_fused_gate_cumsum_matches_unfused( + cu_seqlens: list[int], + dtype: torch.dtype, + lower_bound: float | None, +): + H, D = 8, 64 + T = cu_seqlens[-1] + N = len(cu_seqlens) - 1 + torch.manual_seed(123) + + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE) + q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE) + beta_storage = torch.randn(1, T, 2 * H + 3, dtype=dtype, device=DEVICE) + raw_beta = beta_storage[..., 1 : 2 * H + 1 : 2] + beta = raw_beta.float().sigmoid() + A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous() + dt_bias = ( + torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1 + ).contiguous() + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE) + initial_state = h0.transpose(-1, -2).contiguous() + + gate = fused_kda_gate( + raw_g.reshape(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ) + if lower_bound is not None: + expected_gate = lower_bound * torch.sigmoid( + A_log.exp()[None, :, None] + * (raw_g.float().view(T, H, D) + dt_bias.view(H, D)) + ) + torch.testing.assert_close(gate, expected_gate) + gate = gate.unsqueeze(0) + old_o, old_ht = chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=gate, + beta=beta, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + new_o, new_ht = chunk_kda_with_fused_gate( + q=q.clone(), + k=k.clone(), + v=v.clone(), + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state.clone(), + output_final_state=True, + cu_seqlens=cu_seqlens_t, + use_qk_l2norm_in_kernel=True, + ) + + assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) + assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize("num_seqs", [1, 8, 32]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@pytest.mark.parametrize("state_indices_stride", [1, 8]) +@torch.inference_mode() +def test_packed_kda_decode_correctness( + num_seqs: int, + lower_bound: float | None, + state_indices_stride: int, +): + H, D = 8, 128 + torch.manual_seed(321) + + packed_storage = torch.randn( + num_seqs, + 3 * H * D + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + mixed_qkv = packed_storage[:, : 3 * H * D] + assert mixed_qkv.stride(0) == 3 * H * D + 1 + q, k, v = ( + x.contiguous().view(1, num_seqs, H, D) for x in mixed_qkv.split(H * D, dim=-1) + ) + raw_g = torch.randn( + 1, + num_seqs, + H, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = torch.randn( + 1, + num_seqs, + H, + dtype=torch.bfloat16, + device=DEVICE, + ) + beta = raw_beta.float().sigmoid() + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + state_storage = torch.randn( + num_seqs + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(num_seqs + 1, H, D, D) + assert not state.is_contiguous() + assert state.stride()[1:] == (D * D, D, 1) + state_indices_storage = torch.zeros( + num_seqs, + state_indices_stride, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = state_indices_storage[:, 0] + state_indices.copy_( + torch.arange( + 1, + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ) + ) + gate = fused_kda_gate( + raw_g.reshape(num_seqs, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + dense_state = state.clone() + dense_out, _ = fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=D**-0.5, + initial_state=dense_state, + inplace_final_state=True, + cu_seqlens=torch.arange( + num_seqs + 1, + dtype=torch.int32, + device=DEVICE, + ), + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + ) + packed_state = state + packed_out, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=packed_state, + state_indices=state_indices, + ) + + assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3) + assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3) + + +@pytest.mark.parametrize( + ("H", "fuse_gate"), + [(12, True), (12, False), (12, None), (96, None)], +) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@torch.inference_mode() +def test_kda_spec_decode_correctness( + H: int, + fuse_gate: bool | None, + lower_bound: float | None, +): + num_seqs, query_len, D = 3, 3, 128 + T = num_seqs * query_len + torch.manual_seed(1234) + + qkv_storage = torch.randn( + 1, + T, + 3 * H * D + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + packed_qkv = qkv_storage[..., : 3 * H * D] + q, k, v = (x.view(1, T, H, D) for x in packed_qkv.split(H * D, dim=-1)) + gate_storage = torch.randn( + 1, + T, + H * D + 5, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_g = gate_storage[..., : H * D].view(1, T, H, D) + beta_storage = torch.randn( + 1, + T, + H + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = beta_storage[..., :H] + A_log = 0.5 * torch.randn(H, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(H, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.arange( + 0, + T + 1, + query_len, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = torch.arange( + 1, + T + 1, + dtype=torch.int32, + device=DEVICE, + ).view(num_seqs, query_len) + num_accepted_tokens = torch.tensor( + [1, 2, 3], + dtype=torch.int32, + device=DEVICE, + ) + state_storage = 0.01 * torch.randn( + T + 1, + H * D * D + 17, + dtype=torch.float32, + device=DEVICE, + ) + state = state_storage[:, : H * D * D].view(T + 1, H, D, D) + output_storage = torch.full( + (1, T, H * D + 11), + torch.nan, + dtype=torch.bfloat16, + device=DEVICE, + ) + output = output_storage[..., : H * D].view(1, T, H, D) + + gate = fused_kda_gate( + raw_g.contiguous().view(T, H * D), + A_log, + D, + g_bias=dt_bias, + lower_bound=lower_bound, + ).unsqueeze(0) + beta = raw_beta.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + expected_state = state.clone() + expected_outputs = [] + for seq, accepted in enumerate(num_accepted_tokens.tolist()): + recurrent_state = expected_state[state_indices[seq, accepted - 1]].transpose( + -1, -2 + ) + start = seq * query_len + for token in range(query_len): + token_slice = slice(start + token, start + token + 1) + token_output, recurrent_state = naive_recurrent_kda( + q_norm[:, token_slice], + k_norm[:, token_slice], + v[:, token_slice], + gate[:, token_slice], + beta[:, token_slice], + initial_state=recurrent_state, + output_final_state=True, + ) + assert recurrent_state is not None + expected_outputs.append(token_output) + expected_state[state_indices[seq, token]] = recurrent_state.transpose( + -1, -2 + ) + expected = torch.cat(expected_outputs, dim=1) + + actual_state = state.clone() + actual, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=actual_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + out=output, + fuse_gate=fuse_gate, + ) + + assert actual.data_ptr() == output.data_ptr() + assert_close("o", expected, actual, 1e-3, err_atol=1e-3) + used_states = state_indices.flatten().long() + assert_close( + "ht", + expected_state[used_states], + actual_state[used_states], + 3e-3, + err_atol=3e-3, + ) + assert torch.isnan(output_storage[..., H * D :]).all() + + +@pytest.mark.parametrize( + ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), + [ + (12, 1, -5.0, True), + (12, 4, None, False), + (24, 4, None, False), + (48, 1, -5.0, True), + (96, 1, -5.0, True), + ], +) +@torch.inference_mode() +def test_fused_kda_decode_correctness( + num_heads: int, + num_seqs: int, + lower_bound: float | None, + fuse_output_norm: bool, +): + D, W = 128, 4 + if not is_fused_kda_decode_supported( + num_heads, + D, + W, + num_spec=0, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ): + pytest.skip("Fused KDA decode is not supported on this platform") + torch.manual_seed(967 + num_heads + num_seqs) + dim = num_heads * D + slots = num_seqs + 2 + packed_x_storage = torch.randn( + num_seqs, 3 * dim + 17, dtype=torch.bfloat16, device=DEVICE + ) + packed_x = packed_x_storage[:, : 3 * dim] + weight = 0.1 * torch.randn(3 * dim, W, dtype=torch.float32, device=DEVICE) + conv_seed = 0.1 * torch.randn( + slots, + W - 1, + 3 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ).transpose(1, 2) + raw_g = torch.randn( + 1, + num_seqs, + num_heads, + D, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta_storage = torch.randn( + 1, + num_seqs, + num_heads + 1, + dtype=torch.bfloat16, + device=DEVICE, + ) + raw_beta = raw_beta_storage[:, :, :num_heads] + output_gate_storage = torch.randn( + num_seqs, + dim + 7, + dtype=torch.bfloat16, + device=DEVICE, + ) + output_gate = output_gate_storage[:, :dim].view(num_seqs, num_heads, D) + norm_weight = torch.randn(D, dtype=torch.float32, device=DEVICE) + norm_eps = 1e-5 + A_log = 0.5 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) + dt_bias = 0.1 * torch.randn(dim, dtype=torch.float32, device=DEVICE) + state_indices = torch.arange( + num_seqs, + 0, + -1, + dtype=torch.int32, + device=DEVICE, + ) + state_seed = 0.01 * torch.randn( + slots, + num_heads, + D, + D, + dtype=torch.float32, + device=DEVICE, + ) + + conv_ref = conv_seed.clone() + state_ref = state_seed.clone() + mixed_qkv = causal_conv1d_update( + packed_x, + conv_ref, + weight, + activation="silu", + conv_state_indices=state_indices, + validate_data=True, + out=torch.empty_like(packed_x), + ) + expected, _ = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=state_ref, + state_indices=state_indices, + ) + if fuse_output_norm: + expected_float = expected.float() + expected = ( + expected_float + * torch.rsqrt(expected_float.square().mean(dim=-1, keepdim=True) + norm_eps) + * norm_weight + * output_gate.float().sigmoid().unsqueeze(0) + ).to(expected.dtype) + + conv_slot_elements = 3 * dim * (W - 1) + state_slot_elements = num_heads * D * D + conv_slot_bytes = conv_slot_elements * torch.bfloat16.itemsize + page_bytes = conv_slot_bytes + state_slot_elements * torch.float32.itemsize + cache_storage = torch.empty(slots * page_bytes, dtype=torch.uint8, device=DEVICE) + conv_actual = torch.as_strided( + cache_storage.view(torch.bfloat16), + size=(slots, 3 * dim, W - 1), + stride=(page_bytes // torch.bfloat16.itemsize, 1, 3 * dim), + ) + state_actual = torch.as_strided( + cache_storage.view(torch.float32), + size=(slots, num_heads, D, D), + stride=(page_bytes // torch.float32.itemsize, D * D, D, 1), + storage_offset=conv_slot_bytes // torch.float32.itemsize, + ) + conv_actual.copy_(conv_seed) + state_actual.copy_(state_seed) + fused_weight = weight.reshape(3, dim, W).transpose(1, 2).contiguous() + actual = ops.fused_kda_decode( + x=packed_x, + weight=fused_weight, + bias=None, + conv_state=conv_actual, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + state=state_actual, + lower_bound=lower_bound, + output_gate=output_gate if fuse_output_norm else None, + norm_weight=norm_weight if fuse_output_norm else None, + norm_eps=norm_eps, + ) + + torch.testing.assert_close(actual, expected, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(conv_actual, conv_ref, atol=0, rtol=0) + torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) + + +def test_fused_kda_decode_rejects_speculative_conv_state(): + assert not is_fused_kda_decode_supported( + num_heads=12, + head_dim=128, + conv_width=4, + num_spec=2, + input_dtype=torch.bfloat16, + conv_state_dtype=torch.bfloat16, + ) + + +@torch.inference_mode() +def test_flashkda_correctness(): + if not is_flashkda_supported(128, torch.bfloat16, -3.0): + pytest.skip("FlashKDA is not supported on this platform") + + import vllm._flashkda_C # noqa: F401 + + B, T, H, D = 1, 48, 2, 128 + torch.manual_seed(11) + q, k, v, raw_g = [ + torch.randn(B, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) + ] + beta_logits = torch.randn(B, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = torch.randn(2, H, D, D, dtype=torch.float32, device=DEVICE) + cu_seqlens = torch.tensor([0, 17, T], dtype=torch.int32, device=DEVICE) + lower_bound = -3.0 + + gate = lower_bound * torch.sigmoid( + A_log.exp()[None, None, :, None] * (raw_g.float() + dt_bias[None, None, :, :]) + ) + beta = beta_logits.float().sigmoid() + q_norm = l2norm_fwd(q.contiguous()) + k_norm = l2norm_fwd(k.contiguous()) + + expected_outputs = [] + expected_states = [] + for i, (start, end) in enumerate( + zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()) + ): + output, final_state = naive_recurrent_kda( + q_norm[:, start:end], + k_norm[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=initial_state[i].transpose(-1, -2), + output_final_state=True, + ) + expected_outputs.append(output) + expected_states.append(final_state) + expected_out = torch.cat(expected_outputs, dim=1) + expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous() + + actual_out = torch.empty_like(v) + actual_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size(T, H, cu_seqlens.numel() - 1), + dtype=torch.uint8, + device=DEVICE, + ) + torch.ops._flashkda_C.fwd( + q, + k, + v, + raw_g, + beta_logits, + D**-0.5, + actual_out, + workspace, + A_log, + dt_bias, + lower_bound, + initial_state, + actual_state, + cu_seqlens, + ) + + assert_close("o", expected_out, actual_out, 0.01) + assert_close("ht", expected_state, actual_state, 0.01) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py new file mode 100644 index 00000000000..5352ef7a7a6 --- /dev/null +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import fields + +import pytest +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, + KimiK3KDAMetadataBuilder, + _mamba_get_block_table_tensor, + stage_spec_decode_metadata, +) +from vllm.v1.attention.backend import AttentionMetadataBuilder +from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + mamba_get_block_table_tensor, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +PRUNED_METADATA_FIELDS = { + "chunk_indices", + "chunk_offsets", + "prefill_query_start_loc", + "prefill_state_indices", + "prefill_has_initial_state", + "spec_sequence_masks", +} + + +def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata): + for field in fields(KimiK3KDAMetadata): + actual_value = getattr(actual, field.name) + expected_value = getattr(reference, field.name) + if field.name in PRUNED_METADATA_FIELDS: + assert actual_value is None + continue + if ( + field.name in {"spec_token_indx", "non_spec_token_indx"} + and actual.num_spec_decodes > 0 + and actual.num_prefills == 0 + and actual.num_decodes == 0 + ): + assert actual_value is None + continue + if isinstance(actual_value, torch.Tensor): + torch.testing.assert_close(actual_value, expected_value) + elif field.name == "nums_dict": + assert (actual_value is None) == (expected_value is None) + if actual_value is not None: + assert actual_value[8]["tot"] == expected_value[8]["tot"] + torch.testing.assert_close( + actual_value[8]["nums"], expected_value[8]["nums"] + ) + else: + assert actual_value == expected_value + + +def _make_builder( + builder_cls: type[AttentionMetadataBuilder], + num_speculative_tokens: int, + full_cuda_graph: bool, + device: torch.device = DEVICE, + mamba_cache_mode: str = "none", +) -> AttentionMetadataBuilder: + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + if num_speculative_tokens: + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=num_speculative_tokens, + ) + vllm_config.compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE + ) + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + return builder_cls( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=num_speculative_tokens, + ), + layer_names=["layer.0"], + vllm_config=vllm_config, + device=device, + ) + + +@pytest.mark.parametrize( + ( + "batch", + "num_decode_draft_tokens", + "num_speculative_tokens", + "full_cuda_graph", + "is_prefilling", + ), + [ + pytest.param( + BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]), + [2, 2], + 2, + False, + [False, False], + id="pure-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[100, 65, 20], query_lens=[50, 1, 3]), + [-1, -1, 2], + 2, + False, + [True, False, False], + id="mixed-prefill-and-spec-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + None, + 0, + False, + [False, False], + id="regular-decode", + ), + pytest.param( + BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]), + [0, 0], + 2, + False, + [False, False], + id="no-scheduled-draft-tokens", + ), + ], +) +def test_kimi_k3_kda_metadata_matches_shared_gdn( + batch: BatchSpec, + num_decode_draft_tokens: list[int] | None, + num_speculative_tokens: int, + full_cuda_graph: bool, + is_prefilling: list[bool], +): + kwargs: dict[str, torch.Tensor] = {} + if num_decode_draft_tokens is not None: + kwargs = { + "num_decode_draft_tokens_cpu": torch.tensor( + num_decode_draft_tokens, dtype=torch.int32 + ), + "num_accepted_tokens": torch.ones( + batch.batch_size, dtype=torch.int32, device=DEVICE + ), + } + + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build( + 0, + common_attn_metadata, + **kwargs, + ) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens, + full_cuda_graph, + ).build(0, common_attn_metadata, **kwargs) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata(): + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, False, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + # The K3 layer dispatches the non-spec subgroup to packed decode whenever + # it contains no prefill request. + assert actual.num_decodes == 2 + assert actual.num_decode_tokens == 2 + assert actual.num_prefills == 0 + assert actual.num_prefill_tokens == 0 + assert actual.has_initial_state is None + assert actual.nums_dict is None + assert actual.non_spec_query_start_loc is None + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4])) + torch.testing.assert_close( + actual.spec_query_start_loc, + torch.tensor([0, 3], dtype=torch.int32), + ) + + +def test_mixed_regular_and_spec_decode_excludes_request_padding(): + batch = BatchSpec(seq_lens=[16, 65, 20], query_lens=[0, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, False, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(3, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_decodes == 1 + assert actual.non_spec_state_indices_tensor is not None + assert actual.non_spec_state_indices_tensor.shape == (1,) + torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0])) + torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) + + +@pytest.mark.parametrize( + ("seq_len", "expected_has_initial_state"), + [ + pytest.param(1, False, id="first-token-prefill"), + pytest.param(65, True, id="final-one-token-prefill-chunk"), + ], +) +def test_mixed_one_token_prefill_and_spec_decode_uses_prefill_metadata( + seq_len: int, + expected_has_initial_state: bool, +): + batch = BatchSpec(seq_lens=[seq_len, 20], query_lens=[1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([True, False])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, 2], dtype=torch.int32), + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device=DEVICE), + ) + + assert actual.num_prefills == 1 + assert actual.num_prefill_tokens == 1 + assert actual.num_decodes == 0 + assert actual.num_decode_tokens == 0 + assert actual.has_initial_state is not None + assert actual.has_initial_state.tolist() == [expected_has_initial_state] + assert actual.non_spec_query_start_loc is not None + torch.testing.assert_close( + actual.non_spec_query_start_loc, + torch.tensor([0, 1], dtype=torch.int32), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([False, False])) + reference = _make_builder( + GDNAttentionMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + ).build_for_cudagraph_capture(common_attn_metadata) + + assert isinstance(actual, KimiK3KDAMetadata) + _assert_matches_shared_gdn(reference, actual) + + +def test_kimi_k3_kda_backend_uses_private_metadata_builder(): + assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder + assert KimiK3KDAAttentionBackend.is_ssm() + assert issubclass(KimiK3KDAAttentionBackend, GDNAttentionBackend) + assert issubclass(KimiK3KDAMetadata, GDNAttentionMetadata) + assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_stage_spec_decode_metadata_matches_pytorch(): + device = torch.device("cuda") + num_spec_decodes = 33 + batch_size = 65 + num_state_slots = 3 + state_indices = torch.arange( + num_spec_decodes * 32, + dtype=torch.int32, + device=device, + ).reshape(num_spec_decodes, 32)[:, :num_state_slots] + query_start_loc = ( + torch.arange(num_spec_decodes + 1, dtype=torch.int32, device=device) + * num_state_slots + ) + num_accepted_tokens = ( + torch.arange(num_spec_decodes, dtype=torch.int32, device=device) + % num_state_slots + + 1 + ) + + staged_state_indices = torch.empty( + (batch_size, num_state_slots), dtype=torch.int32, device=device + ) + staged_query_start_loc = torch.empty( + batch_size + 1, dtype=torch.int32, device=device + ) + staged_num_accepted_tokens = torch.empty( + batch_size, dtype=torch.int32, device=device + ) + stage_spec_decode_metadata( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + num_spec_decodes=num_spec_decodes, + ) + + expected_state_indices = torch.full_like(staged_state_indices, NULL_BLOCK_ID) + expected_state_indices[:num_spec_decodes] = state_indices + expected_query_start_loc = torch.full( + (batch_size + 1,), + query_start_loc[-1], + dtype=torch.int32, + device=device, + ) + expected_query_start_loc[: num_spec_decodes + 1] = query_start_loc + expected_num_accepted_tokens = torch.ones( + batch_size, dtype=torch.int32, device=device + ) + expected_num_accepted_tokens[:num_spec_decodes] = num_accepted_tokens + + torch.testing.assert_close(staged_state_indices, expected_state_indices) + torch.testing.assert_close(staged_query_start_loc, expected_query_start_loc) + torch.testing.assert_close(staged_num_accepted_tokens, expected_num_accepted_tokens) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_aligned_block_table_matches_shared_gdn(): + device = torch.device("cuda") + seq_lens = torch.tensor( + [0, 1, 15, 16, 17, 31, 32, 33, 511, 512, 513], + dtype=torch.int32, + device=device, + ).repeat(6)[:65] + block_table_storage = torch.arange( + seq_lens.numel() * 128, + dtype=torch.int32, + device=device, + ).reshape(seq_lens.numel(), 128) + block_table = block_table_storage[:, ::2] + kv_cache_spec = MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=2, + ) + + expected = mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + actual = _mamba_get_block_table_tensor( + block_table, + seq_lens, + kv_cache_spec, + "align", + ) + + torch.testing.assert_close(actual, expected) diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py new file mode 100644 index 00000000000..f18f80a36c0 --- /dev/null +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import ray +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from tests.utils import ( + init_test_distributed_environment, + multi_gpu_test, + multi_process_parallel, +) +from vllm.distributed import get_tp_group +from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup +from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import KimiK3LatentMoETailOp +from vllm.platforms import current_platform + +HIDDEN_SIZE = 7168 +LATENT_SIZE = 3584 +EPS = 0.1 + + +@ray.remote(num_gpus=1, max_calls=1) +def _test_latent_moe_tail_worker( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, + pp_size: int, + rank: int, + distributed_init_port: str, +) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + tp_size, + pp_size, + rank, + distributed_init_port, + ) + + torch.manual_seed(0) + rms_weight = 1 + 0.1 * torch.randn( + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + up_weight = ( + torch.randn( + HIDDEN_SIZE, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ) + / LATENT_SIZE**0.5 + ) + + group = get_tp_group().device_group + op = KimiK3LatentMoETailOp.initialize( + hidden_size=HIDDEN_SIZE, + latent_size=LATENT_SIZE, + dtype=torch.bfloat16, + device=device, + rms_eps=EPS, + ) + cutedsl_warmup() + + for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)): + torch.manual_seed(100 * iteration + rank + 1) + routed_output = torch.randn( + num_tokens, + LATENT_SIZE, + device=device, + dtype=torch.bfloat16, + ).mul_(0.01) + shared_output = torch.randn( + num_tokens, + HIDDEN_SIZE, + device=device, + dtype=torch.bfloat16, + ) + + routed_reference = routed_output.clone() + shared_reference = shared_output.clone() + dist.all_reduce(routed_reference, group=group) + dist.all_reduce(shared_reference, group=group) + expected = F.linear( + F.rms_norm( + routed_reference, + (LATENT_SIZE,), + rms_weight, + EPS, + ), + up_weight, + ) + expected.add_(shared_reference) + + actual = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) + assert actual.is_contiguous() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + graph.replay() + torch.testing.assert_close(graph_output, expected, atol=8e-2, rtol=3e-2) + + +def _run_latent_moe_tail_test( + monkeypatch: pytest.MonkeyPatch, + tp_size: int, +) -> None: + if not current_platform.is_device_capability_family(100): + pytest.skip("K3 latent-MoE tail fusion requires SM100") + multi_process_parallel( + monkeypatch, + tp_size, + 1, + _test_latent_moe_tail_worker, + ) + + +@multi_gpu_test(num_gpus=8) +def test_latent_moe_tail_tp8_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 8) + + +@multi_gpu_test(num_gpus=16) +def test_latent_moe_tail_tp16_matches_native_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run_latent_moe_tail_test(monkeypatch, 16) diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py new file mode 100644 index 00000000000..c65f9bd9532 --- /dev/null +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import MethodType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +from vllm.config import ParallelConfig +from vllm.models.kimi_k3.nvidia import model as kimi_model +from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp +from vllm.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops +from vllm.platforms import current_platform + + +class _IdentityNorm(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False) + self.variance_epsilon = 1e-5 + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + ): + if residual is None: + return hidden_states + return hidden_states, residual + + +class _RecordingMoE(nn.Module): + def __init__(self) -> None: + super().__init__() + self.num_tokens = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.num_tokens = hidden_states.shape[0] + return hidden_states + + +class _Projection(nn.Module): + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.ones(1, hidden_size), + requires_grad=False, + ) + + +class _SequenceParallelMTPBlock: + use_sequence_parallel = True + + def __call__( + self, + *, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ): + assert residual is None + return hidden_states * 2, None, hidden_states * 3 + + +def _mock_sequence_parallel_collectives(monkeypatch): + monkeypatch.setattr( + kimi_model, + "sp_reduce_scatter", + lambda tensor: tensor.chunk(2, dim=0)[0], + ) + monkeypatch.setattr( + kimi_model, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_model, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + +@pytest.mark.parametrize( + ("num_tokens", "is_padding", "tp_rank", "expected"), + [ + (1, None, 0, [False]), + (1, None, 1, [True]), + (5, None, 2, [False, True]), + (5, None, 3, [True, True]), + (5, [False, True, False, False, False], 0, [False, True]), + ], +) +def test_sp_padding_mask_marks_added_rows( + monkeypatch, + num_tokens: int, + is_padding: list[bool] | None, + tp_rank: int, + expected: list[bool], +): + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) + + hidden_states = torch.empty(num_tokens, 2) + padding = torch.tensor(is_padding) if is_padding is not None else None + actual = sp_ops.sp_padding_mask(padding, hidden_states) + + torch.testing.assert_close(actual, torch.tensor(expected)) + + +@pytest.mark.parametrize( + ("data_parallel_size", "expected"), + [ + (1, False), + (2, True), + ], +) +def test_moe_sequence_parallel_requires_data_parallel( + monkeypatch, + data_parallel_size: int, + expected: bool, +): + monkeypatch.setattr(current_platform, "device_count", lambda: 2) + parallel_config = ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=data_parallel_size, + enable_expert_parallel=True, + all2all_backend="allgather_reducescatter", + ) + + assert parallel_config.use_sequence_parallel_moe is expected + + +def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = False + layer.use_sequence_parallel = True + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + + positions = torch.arange(3) + full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + hidden_states = kimi_model.sp_shard(full_hidden_states) + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + assert prefix_sum is None + assert hidden_states.shape == residual.shape == (2, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch): + layer = object.__new__(kimi_model.KimiDecoderLayer) + nn.Module.__init__(layer) + layer.use_attn_res = True + layer.use_sequence_parallel = True + layer.prev_valid_blocks = 0 + layer.block_write_idx = 0 + layer.is_block_write_layer = False + layer.input_layernorm = _IdentityNorm() + layer.post_attention_layernorm = _IdentityNorm() + layer.self_attention_res_norm = _IdentityNorm() + layer.mlp_res_norm = _IdentityNorm() + layer.self_attention_res_proj = _Projection() + layer.mlp_res_proj = _Projection() + layer.mlp = _RecordingMoE() + layer._run_self_attn = MethodType( + lambda self, positions, hidden_states: hidden_states, + layer, + ) + + _mock_sequence_parallel_collectives(monkeypatch) + monkeypatch.setattr( + kimi_model, + "attn_res", + lambda prefix_sum, hidden_states, *args, **kwargs: ( + prefix_sum if hidden_states is None else prefix_sum + hidden_states + ), + ) + + prefix_sum = kimi_model.sp_shard(torch.arange(6, dtype=torch.float32).view(3, 2)) + block_residual = torch.zeros(2, 1, 2) + hidden_states, prefix_sum, block_residual = layer( + positions=torch.arange(3), + hidden_states=None, + prefix_sum=prefix_sum, + residual=block_residual, + ) + + assert hidden_states.shape == prefix_sum.shape == (2, 2) + assert block_residual.shape == (2, 1, 2) + assert layer.mlp.num_tokens == 2 + + +def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): + layer = object.__new__(kimi_mtp.KimiK3MultiTokenPredictorLayer) + nn.Module.__init__(layer) + layer.enorm = _IdentityNorm() + layer.hnorm = _IdentityNorm() + layer.eh_proj = nn.Identity() + object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock()) + + final_norm = Mock(side_effect=lambda hidden_states: hidden_states + 1) + object.__setattr__( + layer, + "shared_head", + SimpleNamespace(norm=final_norm), + ) + + monkeypatch.setattr( + kimi_mtp, + "fused_mtp_input", + lambda positions, inputs_embeds, *args: inputs_embeds, + ) + monkeypatch.setattr( + kimi_mtp, + "sp_shard", + lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2], + ) + monkeypatch.setattr( + kimi_mtp, + "sp_all_gather", + lambda tensor: torch.cat([tensor, tensor], dim=0), + ) + + inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2) + logits_hidden_states, hidden_states = layer( + input_ids=torch.zeros(3, dtype=torch.long), + positions=torch.arange(3), + previous_hidden_states=torch.zeros_like(inputs_embeds), + inputs_embeds=inputs_embeds, + ) + + sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2] + expected_hidden_states = torch.cat( + [sharded_states * 5, sharded_states * 5], + dim=0, + )[:3] + torch.testing.assert_close(hidden_states, expected_hidden_states) + torch.testing.assert_close(logits_hidden_states, expected_hidden_states + 1) + final_norm.assert_called_once() + torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states) + + +def test_sp_all_gather_uses_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + expected = torch.cat([hidden_states, hidden_states]) + custom_all_gather = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_all_gather=custom_all_gather, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", fallback) + + output = sp_ops.sp_all_gather(hidden_states) + + torch.testing.assert_close(output, expected) + custom_all_gather.assert_called_once_with(hidden_states) + fallback.assert_not_called() + + +def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch): + hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + custom_reduce_scatter = Mock(return_value=expected) + device_communicator = SimpleNamespace( + custom_reduce_scatter=custom_reduce_scatter, + ) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=device_communicator), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + fallback = Mock(side_effect=AssertionError("unexpected fallback")) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_reduce_scatter", fallback) + + output = sp_ops.sp_reduce_scatter(hidden_states) + + torch.testing.assert_close(output, expected) + padded = custom_reduce_scatter.call_args.args[0] + assert padded.shape == (4, 2) + torch.testing.assert_close(padded[:3], hidden_states) + torch.testing.assert_close(padded[3], torch.zeros(2)) + fallback.assert_not_called() + + +def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch): + hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2) + monkeypatch.setattr( + sp_ops, + "get_tp_group", + lambda: SimpleNamespace(device_communicator=None), + ) + monkeypatch.setattr( + sp_ops, + "get_tensor_model_parallel_world_size", + lambda: 2, + ) + all_gather = Mock(return_value=hidden_states) + reduce_scatter = Mock(return_value=hidden_states) + monkeypatch.setattr(sp_ops, "tensor_model_parallel_all_gather", all_gather) + monkeypatch.setattr( + sp_ops, + "tensor_model_parallel_reduce_scatter", + reduce_scatter, + ) + + torch.testing.assert_close(sp_ops.sp_all_gather(hidden_states), hidden_states) + torch.testing.assert_close(sp_ops.sp_reduce_scatter(hidden_states), hidden_states) + all_gather.assert_called_once_with(hidden_states, 0) + reduce_scatter.assert_called_once_with(hidden_states, 0) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index ecf1ad26d0c..5c431140b04 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -1242,7 +1242,7 @@ def test_custom_inputs_models( create_new_process_for_each_test=True, ), ) -@create_new_process_for_each_test() +@create_new_process_for_each_test("spawn") def test_single_image_models_heavy( tmp_path: PosixPath, model_type: str, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 9f592386c15..bc6d6d4afa4 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2733,6 +2733,52 @@ def fused_minimax_m3_qknorm_rope_kv_insert( ) +def fused_kda_decode( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_state: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_indices: torch.Tensor, + state: torch.Tensor, + out: torch.Tensor | None = None, + lower_bound: float | None = None, + output_gate: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-5, +) -> torch.Tensor: + if out is None: + out = torch.empty( + 1, + x.shape[0], + raw_g.shape[2], + raw_g.shape[3], + dtype=x.dtype, + device=x.device, + ) + torch.ops._C.fused_kda_decode( + x, + weight, + bias, + conv_state, + raw_g, + raw_beta, + A_log, + dt_bias, + state_indices, + state, + out, + lower_bound, + output_gate, + norm_weight, + norm_eps, + ) + return out + + def concat_and_cache_mla( kv_c: torch.Tensor, k_pe: torch.Tensor, @@ -2746,6 +2792,26 @@ def concat_and_cache_mla( ) +def concat_and_cache_mla_grouped( + kv_c: torch.Tensor, + k_pe: torch.Tensor, + kv_cache_ptrs: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int, + block_stride: int, + entry_stride: int, +) -> None: + torch.ops._C_cache_ops.concat_and_cache_mla_grouped( + kv_c, + k_pe, + kv_cache_ptrs, + slot_mapping, + block_size, + block_stride, + entry_stride, + ) + + def kimi_k3_attn_res( prefix: torch.Tensor, delta: torch.Tensor, @@ -3040,6 +3106,63 @@ def all_reduce( torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) +def custom_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_all_gather( + fa, inp, out, reg_buffer, reg_buffer_sz_bytes + ) + + +def mnnvl_lamport_all_gather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + multicast_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_all_gather( + fa, + inp, + out, + local_buffer, + multicast_buffer, + epoch_buffer, + stage_sz_bytes, + ) + + +def custom_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.custom_reduce_scatter( + fa, inp, out, reg_buffer, reg_buffer_sz_bytes + ) + + +def mnnvl_lamport_reduce_scatter( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + local_buffer: int, + epoch_buffer: int, + stage_sz_bytes: int, +) -> None: + torch.ops._C_custom_ar.mnnvl_lamport_reduce_scatter( + fa, inp, out, local_buffer, epoch_buffer, stage_sz_bytes + ) + + def dispose(fa: int) -> None: torch.ops._C_custom_ar.dispose(fa) @@ -3144,18 +3267,18 @@ def dsv3_fused_a_gemm( output: torch.Tensor, mat_a: torch.Tensor, mat_b: torch.Tensor, + enable_pdl: bool = False, ) -> None: - """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). + """Low-latency fused-A-style GEMM (SM 9.0+, BF16, 1-16 tokens). - Computes output = mat_a @ mat_b.T where: - mat_a: [num_tokens, 7168] row-major bf16 (hidden states) - mat_b: [7168, 2112] column-major bf16 (weight transposed) - output: [num_tokens, 2112] row-major bf16 + Computes ``output = mat_a @ mat_b`` for the compiled Kimi K3 and + DeepSeek V3 projection shapes. ``mat_a`` and ``output`` are row-major; + ``mat_b`` is the column-major transposed weight. ``enable_pdl`` permits + programmatic dependent launch for callers that have validated it. - Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes. - Requires SM 9.0+ (Hopper). + Requires SM 9.0+. """ - torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b) + torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl) if hasattr(torch.ops._C, "weight_packed_linear"): diff --git a/vllm/envs.py b/vllm/envs.py index 84ec5a8af85..607b6ce25dd 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -269,6 +269,7 @@ if TYPE_CHECKING: VLLM_NCCL_INCLUDE_PATH: str | None = None VLLM_GC_DEBUG: str = "" VLLM_DEBUG_WORKSPACE: bool = False + VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 @@ -1894,6 +1895,11 @@ environment_variables: dict[str, Callable[[], Any]] = { # Debug workspace allocations. # logging of workspace resize operations. "VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))), + # Enable the experimental Kimi K3 latent-MoE tail fusion. + # Currently supported only on SM100 with TP=8/16 and BF16. + "VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION": lambda: bool( + int(os.getenv("VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION", "0")) + ), # Disables parallel execution of shared_experts via separate cuda stream "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool( int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0")) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py b/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py new file mode 100644 index 00000000000..44964f3e47b --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_skinny_gemm.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr + + +class CuteSkinnyGemm: + """Shape-dynamic low-latency GEMM for small token counts. + + Computes ``C[M, N] = A[M, K] @ B[N, K].T + residual`` with BF16 or FP16 + inputs, FP32 accumulators, and an output matching the input dtype. The + residual term is optional and is added before the output conversion. N and + K are runtime values. The tiny M dimension is fully unrolled alongside a + small set of tuning parameters. + """ + + def __init__( + self, + *, + element_type, + num_rows: int, + block_size: int, + outputs_per_block: int, + vector_width: int = 8, + k_unroll: int = 1, + has_residual: bool = False, + use_pdl: bool = False, + ) -> None: + if block_size % cute.arch.WARP_SIZE != 0: + raise ValueError("block_size must be a multiple of the warp size") + self.element_type = element_type + self.num_rows = num_rows + self.block_size = block_size + self.outputs_per_block = outputs_per_block + self.vector_width = vector_width + self.k_unroll = k_unroll + self.has_residual = has_residual + self.use_pdl = use_pdl + self.num_warps = block_size // cute.arch.WARP_SIZE + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + stream: CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(gA, gB, gResidual, gC, k, copy_a, copy_b).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=self.num_rows * self.outputs_per_block * self.num_warps * 4, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gResidual: cute.Tensor, + gC: cute.Tensor, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + num_rows: cutlass.Constexpr = self.num_rows + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + + acc_layout = cute.make_layout( + (num_rows, outputs_per_block), stride=(outputs_per_block, 1) + ) + acc = cute.make_rmem_tensor(acc_layout, cutlass.Float32) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + # Layout after both divides is (M/N, K_TILE, K_LANE, K_VEC). + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout((num_rows, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, vector_width), stride=(vector_width, 1) + ), + self.element_type, + ) + + for k_tile in cutlass.range(num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, k_tile], a_regs[mi, None]) + + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + tB = tB_all[n_idx, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to( + cutlass.Float32 + ) * b_regs[ni, vi].to(cutlass.Float32) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + n_idx = n_base + ni + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ) + if const_expr(self.has_residual): + total += gResidual[mi, n_idx].to(cutlass.Float32) + gC[mi, n_idx] = cutlass.Float32(total).to(self.element_type) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py new file mode 100644 index 00000000000..e1dd18920b7 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/skinny_gemm.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_cutedsl_available: bool | None = None + + +@dataclass(frozen=True, slots=True) +class SkinnyGemmConfig: + num_rows: int + block_size: int + outputs_per_block: int + k_unroll: int = 1 + vector_width: int = 8 + + +class ShapeDynamicSkinnyGemm: + def __init__(self) -> None: + self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {} + self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set() + self._warmup_registered = False + + @staticmethod + def is_available() -> bool: + global _cutedsl_available + if _cutedsl_available is not None: + return _cutedsl_available + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + + _cutedsl_available = True + except ImportError: + _cutedsl_available = False + logger.info("cuteDSL is not available; skinny GEMM is disabled") + return _cutedsl_available + + @staticmethod + def _config(m: int, n: int, k: int) -> SkinnyGemmConfig: + num_rows = m + wide_block = 224 + if m == 1 and k >= 7168 and k % (wide_block * 8) == 0: + if n % 3 == 0: + k_unroll = 2 if n <= 2304 else 4 + return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll) + if 2304 < n < 4096 and n % 2 == 0: + return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4) + + if k <= 2048 or k % (128 * 8) != 0: + outputs_per_block = 2 if k <= 2048 else 4 + if n % outputs_per_block: + outputs_per_block = 1 + if k % (64 * 8) == 0: + return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2) + if k % (32 * 8) == 0: + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2) + return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4) + + block_size = 64 if 4096 <= n < 8192 else 128 + outputs_per_block = 1 if m == 1 and n <= 2304 else 2 + if n % outputs_per_block: + outputs_per_block = 1 + k_unroll = 2 if n <= 2304 or n >= 16384 else 1 + return SkinnyGemmConfig( + num_rows, + block_size, + outputs_per_block, + k_unroll=k_unroll, + ) + + @staticmethod + def _cutlass_dtype(dtype: torch.dtype): + from cutlass import BFloat16, Float16 + + return BFloat16 if dtype == torch.bfloat16 else Float16 + + @staticmethod + def _stream(): + from cuda.bindings.driver import CUstream + + from vllm.utils.torch_utils import current_stream + + return CUstream(current_stream().cuda_stream) + + @staticmethod + def _use_pdl() -> bool: + from vllm.platforms import current_platform + + return current_platform.is_arch_support_pdl() + + def _compile( + self, + dtype: torch.dtype, + config: SkinnyGemmConfig, + has_residual: bool, + ) -> None: + import cutlass.cute as cute + from quack.compile_utils import make_fake_tensor + + from ._skinny_gemm import CuteSkinnyGemm + + element_type = self._cutlass_dtype(dtype) + n = cute.sym_int(divisibility=config.outputs_per_block) + k = cute.sym_int(divisibility=config.block_size * config.vector_width) + a = make_fake_tensor( + element_type, + (config.num_rows, k), + divisibility=config.vector_width, + ) + b = make_fake_tensor( + element_type, + (n, k), + divisibility=config.vector_width, + ) + c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) + kernel = CuteSkinnyGemm( + element_type=element_type, + num_rows=config.num_rows, + block_size=config.block_size, + outputs_per_block=config.outputs_per_block, + vector_width=config.vector_width, + k_unroll=config.k_unroll, + has_residual=has_residual, + use_pdl=self._use_pdl(), + ) + self._compiled[(dtype, config, has_residual)] = cute.compile( + kernel, + a, + b, + residual, + c, + self._stream(), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", + ) + + def request_warmup( + self, + dtype: torch.dtype, + shapes: Iterable[tuple[int, int, int]], + ) -> None: + """Request compilation of ``(M, N, K)`` shapes before graph capture.""" + self.request_warmup_configs( + dtype, + (self._config(m, n, k) for m, n, k in shapes), + ) + + def request_warmup_configs( + self, + dtype: torch.dtype, + configs: Iterable[SkinnyGemmConfig], + *, + has_residual: bool = False, + ) -> None: + """Request compilation of explicit measured configs before capture.""" + self._warmup_configs.update((dtype, config, has_residual) for config in configs) + if self._warmup_registered: + return + from vllm.model_executor.warmup.cutedsl_warmup import ( + register_cutedsl_warmup_provider, + ) + + register_cutedsl_warmup_provider(self) + self._warmup_registered = True + + def get_cutedsl_warmup_compile_units(self): + from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit + + return tuple( + CuTeDSLCompileUnit( + name=( + "shape-dynamic skinny GEMM with residual" + if has_residual + else "shape-dynamic skinny GEMM" + ), + key=("shape-dynamic-skinny-gemm", dtype, config, has_residual), + compile=partial(self._compile, dtype, config, has_residual), + ) + for dtype, config, has_residual in sorted( + self._warmup_configs, + key=lambda item: ( + str(item[0]), + item[1].num_rows, + item[1].block_size, + item[1].outputs_per_block, + item[1].k_unroll, + item[1].vector_width, + item[2], + ), + ) + ) + + def __call__( + self, + a: torch.Tensor, + b: torch.Tensor, + config: SkinnyGemmConfig | None = None, + residual: torch.Tensor | None = None, + ) -> torch.Tensor: + if a.dim() != 2 or b.dim() != 2: + raise ValueError("a and b must be 2D tensors") + if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype: + raise ValueError("a and b must have the same BF16 or FP16 dtype") + if not a.is_cuda or not b.is_cuda or a.device != b.device: + raise ValueError("a and b must be CUDA tensors on the same device") + if not a.is_contiguous() or not b.is_contiguous(): + raise ValueError("a and b must be contiguous") + if a.shape[1] != b.shape[1]: + raise ValueError("a and b must have matching K dimensions") + if not 1 <= a.shape[0] <= 16: + raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16") + if residual is not None: + if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]): + raise ValueError("residual must have shape (M, N)") + if residual.dtype != a.dtype: + raise ValueError("residual must have the same dtype as a and b") + if residual.device != a.device or not residual.is_cuda: + raise ValueError("residual must be on the same CUDA device") + if not residual.is_contiguous(): + raise ValueError("residual must be contiguous") + + config = config or self._config(a.shape[0], b.shape[0], a.shape[1]) + if config.num_rows != a.shape[0]: + raise ValueError("config num_rows must match M") + if b.shape[0] % config.outputs_per_block != 0: + raise ValueError("N must be divisible by outputs_per_block") + if a.shape[1] % (config.block_size * config.vector_width) != 0: + raise ValueError( + "K must be divisible by block_size * vector_width for this config" + ) + has_residual = residual is not None + cache_key = (a.dtype, config, has_residual) + if cache_key not in self._compiled: + self._compile(a.dtype, config, has_residual) + output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device) + residual_arg = output if residual is None else residual + self._compiled[cache_key](a, b, residual_arg, output, self._stream()) + return output + + +shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm() diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index a8c91af9adc..c8a51f4c220 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -158,6 +158,51 @@ class SiluAndMul(CustomOp): return self.forward_native(x) +@CustomOp.register("situ_and_mul") +class SituAndMul(CustomOp): + """SituGLU activation used by Kimi models. + + Computes beta * tanh(gate / beta) * sigmoid(gate) * up. When + ``linear_beta`` is set, the up projection is also softly clipped with + linear_beta * tanh(up / linear_beta). + """ + + def __init__( + self, + beta: float = 1.0, + linear_beta: float | None = None, + *, + compile_native: bool = True, + ): + super().__init__(compile_native=compile_native) + self.beta = float(beta) + self.linear_beta = None if linear_beta is None else float(linear_beta) + if current_platform.is_cuda_alike(): + self.op = torch.ops._C.situ_and_mul + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].float() + up = x[..., d:].float() + gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (gate * up).to(x.dtype) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries. + # linear_beta<=0 signals "unset" to the kernel (up passed through). + d = x.shape[-1] // 2 + out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device) + self.op( + out, x, self.beta, -1.0 if self.linear_beta is None else self.linear_beta + ) + return out + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + @CustomOp.register("silu_and_mul_with_clamp") class SiluAndMulWithClamp(CustomOp): """SwiGLU activation with input clamping (used by some MoE shared experts). diff --git a/vllm/model_executor/layers/fused_moe/activation.py b/vllm/model_executor/layers/fused_moe/activation.py index d114992504b..baa6136485e 100644 --- a/vllm/model_executor/layers/fused_moe/activation.py +++ b/vllm/model_executor/layers/fused_moe/activation.py @@ -22,6 +22,7 @@ class MoEActivation(Enum): # expects the *packed* layout ([all gates; all ups]), as produced by a # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3). SWIGLUOAI = "swigluoai" + SITU = "situ" SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave" SWIGLUSTEP = "swiglustep" @@ -77,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = { MoEActivation.SILU: "silu_and_mul", MoEActivation.GELU: "gelu_and_mul", MoEActivation.GELU_TANH: "gelu_tanh_and_mul", + MoEActivation.SITU: "situ_and_mul", MoEActivation.SWIGLUOAI: "swigluoai_and_mul", MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp", MoEActivation.SWIGLUSTEP: "swiglustep_and_mul", @@ -133,6 +135,8 @@ def apply_moe_activation( beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """Apply MoE activation function. @@ -163,6 +167,25 @@ def apply_moe_activation( torch.ops._C.gelu_and_mul(output, input) elif activation == MoEActivation.GELU_TANH: torch.ops._C.gelu_tanh_and_mul(output, input) + elif activation == MoEActivation.SITU: + # Fused CUDA kernel: writes straight to `output`, no fp32 temporaries. + # (The pure-torch fallback below upcast both halves to fp32 and + # allocated ~8 temporaries per call, blowing up MoE memory.) + # Both betas come from FusedMoEConfig; a missing beta means the caller + # bypassed the config plumbing, so fail rather than silently use 1.0. + # linear_beta is genuinely optional: <= 0 signals "unset" to the kernel + # (up passed through), matching SituAndMul(linear_beta=None). + assert activation_situ_beta is not None, ( + "SITU requires activation_situ_beta from FusedMoEConfig" + ) + torch.ops._C.situ_and_mul( + output, + input, + activation_situ_beta, + -1.0 + if activation_situ_linear_beta is None + else activation_situ_linear_beta, + ) elif activation == MoEActivation.SWIGLUOAI: torch.ops._C.swigluoai_and_mul(output, input) elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 02e6c51116d..b4aed6acf89 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1314,6 +1314,10 @@ class FusedMoEConfig: swiglu_alpha: float | None = None swiglu_beta: float | None = None + # SituGLU parameters used by Kimi sit(u/v2) activations. + activation_situ_beta: float | None = None + activation_situ_linear_beta: float | None = None + max_capture_size: int = 0 # Set by __post_init__ diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 160cff1eeb5..211364d24bc 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -444,7 +444,14 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP] + # SILU has fused gate+mul+quant kernels; SWIGLUSTEP/SITU take the + # general path (activation applied via self.activation, which forwards + # the situ betas, then FP8 requant). + return activation in [ + MoEActivation.SILU, + MoEActivation.SWIGLUSTEP, + MoEActivation.SITU, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -488,16 +495,15 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): M_sum, N = input.size() activation_out_dim = self.adjust_N_for_activation(N, activation) - if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: - assert activation == MoEActivation.SILU - return fused_silu_mul_fp8_quant_packed( - input=input, - output_q=output, - group_size=block_k, - clamp_limit=self.gemm1_clamp_limit, - ) - if activation == MoEActivation.SILU: + # Fused gate+mul+quant kernels for the common SILU case. + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return fused_silu_mul_fp8_quant_packed( + input=input, + output_q=output, + group_size=block_k, + clamp_limit=self.gemm1_clamp_limit, + ) use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return silu_mul_per_token_group_quant_fp8_colmajor( input=input, @@ -506,12 +512,24 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): clamp_limit=self.gemm1_clamp_limit, ) + # General gated activations (SWIGLUSTEP, SITU): apply the activation + # (self.activation forwards the situ betas from moe_config) then + # FP8-requant into the layout DeepGEMM expects for this scale format. act_out = torch.empty( (M_sum, activation_out_dim), dtype=input.dtype, device=input.device ) self.activation(activation, act_out, input) + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + return per_token_group_quant_fp8_packed_for_deepgemm( + act_out, block_k, use_ue8m0=True, out_q=output + ) + use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return per_token_group_quant_fp8( - act_out, block_k, column_major_scales=True, out_q=output + act_out, + block_k, + column_major_scales=True, + out_q=output, + use_ue8m0=use_ue8m0, ) def apply( diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 20576de5a60..6336154a9ea 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -93,6 +93,8 @@ def _fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: assert hidden_states.ndim == 2 M, K = hidden_states.size() @@ -171,6 +173,8 @@ def _fused_marlin_moe( beta=gemm1_beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) if output is None: @@ -256,6 +260,8 @@ def fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function computes a Mixture of Experts (MoE) layer using two sets of @@ -357,6 +363,8 @@ def fused_marlin_moe( num_tokens_post_padded=num_tokens_post_padded, activation=activation, activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, input_global_scale1=input_global_scale1, input_global_scale2=input_global_scale2, global_scale1=global_scale1, @@ -423,6 +431,9 @@ def batched_fused_marlin_moe( clamp_limit: float | None = None, gemm1_alpha: float = 1.0, gemm1_beta: float = 0.0, + activation_func: Callable[..., None] = apply_moe_activation, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> torch.Tensor: """ This function massages the inputs so the batched hidden_states can be @@ -530,6 +541,9 @@ def batched_fused_marlin_moe( quant_type=quant_type, apply_router_weight_on_input=apply_router_weight_on_input, activation=activation, + activation_func=activation_func, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, expert_map=expert_map, block_size_m=block_size_m, sorted_token_ids=sorted_token_ids, @@ -645,6 +659,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH, + MoEActivation.SITU, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, MoEActivation.SWIGLUSTEP, @@ -793,6 +808,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): global_num_experts=global_num_experts, activation=activation, activation_func=self.activation, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=( + self.moe_config.activation_situ_linear_beta + ), moe_sum=self.moe_sum, expert_map=expert_map, output=output, @@ -833,6 +852,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: # act_input = intermediate_cache1 (M*topk, 2N for gated) # act_output = intermediate_cache2 (M*topk, N) @@ -871,6 +892,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, ) lora_state["cache2"] = act_output @@ -918,6 +941,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): global_num_experts=global_num_experts, activation=activation, activation_func=activation_with_lora, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, moe_sum=moe_sum_with_lora, expert_map=expert_map, output=output, @@ -1021,6 +1046,41 @@ class BatchedMarlinExperts(MarlinExpertsBase): apply_router_weight_on_input: bool, ): assert expert_tokens_meta is not None, "Num valid tokens per batch is required" + + def activation_func( + act: MoEActivation, + act_output: torch.Tensor, + act_input: torch.Tensor, + *, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + **kwargs, + ) -> None: + if act != MoEActivation.SITU: + self.activation( + act, + act_output, + act_input, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + **kwargs, + ) + return + + num_experts, max_num_tokens = hidden_states.shape[:2] + beta = activation_situ_beta + linear_beta = activation_situ_linear_beta + assert beta is not None, ( + "SITU requires activation_situ_beta from FusedMoEConfig" + ) + torch.ops._C.masked_situ_and_mul( + act_output.view(num_experts, max_num_tokens, -1), + act_input.view(num_experts, max_num_tokens, -1), + expert_tokens_meta.expert_num_tokens, + beta, + -1.0 if linear_beta is None else linear_beta, + ) + return batched_fused_marlin_moe( hidden_states=hidden_states, expert_num_tokens=expert_tokens_meta.expert_num_tokens, @@ -1051,4 +1111,7 @@ class BatchedMarlinExperts(MarlinExpertsBase): clamp_limit=self.gemm1_clamp_limit, gemm1_alpha=self.gemm1_alpha, gemm1_beta=self.gemm1_beta, + activation_func=activation_func, + activation_situ_beta=self.moe_config.activation_situ_beta, + activation_situ_linear_beta=self.moe_config.activation_situ_linear_beta, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index e00f6901aab..2027ed297c5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -79,6 +79,36 @@ class TrtLlmMxfp4ExpertsBase: else: self.gemm1_clamp_limit = None + # SITU (SituGLU) TRTLLM-Gen kernel computes + # left = alpha * tanh(x0 / alpha) * sigmoid(x0) # gate (x0) + # right = beta * tanh(x1 / beta) # up (x1) + # which matches vLLM's situ_and_mul with (beta, linear_beta), so map + # situ beta -> gatedActAlpha (gemm1_alpha) and situ linear_beta -> + # gatedActBeta (gemm1_beta). Both must be > 0. + if moe_config.activation == MoEActivation.SITU: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, ( + "SITU requires activation_situ_beta > 0" + ) + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 " + "(the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = torch.full( + (self.local_num_experts,), + float(situ_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_beta = torch.full( + (self.local_num_experts,), + float(situ_linear_beta), + dtype=torch.float32, + device=device, + ) + self.gemm1_clamp_limit = None + self.max_capture_size = moe_config.max_capture_size @staticmethod @@ -103,7 +133,19 @@ class TrtLlmMxfp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + return activation in ( + MoEActivation.SWIGLUOAI, + MoEActivation.SILU, + MoEActivation.SITU, + ) + + @staticmethod + def _flashinfer_activation_type(activation: MoEActivation) -> int: + from flashinfer.fused_moe.core import ActivationType + + if activation == MoEActivation.SITU: + return ActivationType.Situ.value + return ActivationType.Swiglu.value @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -142,6 +184,7 @@ class TrtLlmMxfp4ExpertsMonolithic( activation_key: QuantKey | None, ) -> bool: return routing_method in [ + RoutingMethodType.DeepSeekV3, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, ] @@ -192,8 +235,8 @@ class TrtLlmMxfp4ExpertsMonolithic( device=hidden_states.device, ) trtllm_fp4_block_scale_moe( - routing_logits=router_logits.to(torch.bfloat16), - routing_bias=None, + routing_logits=router_logits, + routing_bias=e_score_correction_bias, hidden_states=x_quant, hidden_states_scale=x_scale, gemm1_weights=w1, @@ -210,14 +253,15 @@ class TrtLlmMxfp4ExpertsMonolithic( output2_scale_scalar=None, num_experts=global_num_experts, top_k=self.topk, - n_group=None, - topk_group=None, + n_group=(num_expert_group or 0), + topk_group=(topk_group or 0), intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, - routed_scaling_factor=None, + routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, do_finalize=True, + activation_type=self._flashinfer_activation_type(activation), tune_max_num_tokens=max(self.max_capture_size, 1), output=output, routing_replay_out=routing_replay_out, @@ -338,6 +382,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula "routing_method_type": RoutingMethodType.Renormalize, "do_finalize": True, "enable_pdl": True, + "activation_type": self._flashinfer_activation_type(activation), "output": output, "tune_max_num_tokens": max(self.max_capture_size, 1), } diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index b6c2429139f..9bdff642813 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -110,6 +110,27 @@ class TrtLlmNvFp4ExpertsBase: self.gemm1_alpha = None self.gemm1_beta = None + # SITU (Kimi SituGLU) TRTLLM-Gen kernel computes + # left=alpha*tanh(x0/alpha)*sigmoid(x0), right=beta*tanh(x1/beta), + # matching vLLM's situ_and_mul, so map situ beta -> gatedActAlpha + # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta). + # These operate on the dequantized gate/up, so they are NOT folded by + # g1_alphas in process_weights_after_loading. + self.is_situ = moe_config.activation == MoEActivation.SITU + if self.is_situ: + situ_beta = moe_config.activation_situ_beta + situ_linear_beta = moe_config.activation_situ_linear_beta + assert situ_beta is not None and situ_beta > 0, ( + "SITU requires activation_situ_beta > 0" + ) + assert situ_linear_beta is not None and situ_linear_beta > 0, ( + "TRTLLM SiTuGlu requires activation_situ_linear_beta > 0 " + "(the private cubin has no up-passthrough path)" + ) + self.gemm1_alpha = _per_expert(situ_beta) + self.gemm1_beta = _per_expert(situ_linear_beta) + self.gemm1_clamp_limit = None + logger.debug_once( "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s", moe_config.activation, @@ -142,7 +163,10 @@ class TrtLlmNvFp4ExpertsBase: # (via the in-place mul above) and never changes again, so this is a # static, per-expert constant. Register on the layer so EPLB # rearranges it alongside the other expert tensors. - if self.gemm1_clamp_limit is not None: + # SITU alpha/beta act on the dequantized gate/up (tanh clamps), not the + # raw GEMM1 accumulator, so they are registered as-is without the + # g1_alphas fold used by the SwiGLU-OAI clamp/beta below. + if self.gemm1_clamp_limit is not None and not self.is_situ: gemm1_clamp_limit = self.gemm1_clamp_limit / self.quant_config.g1_alphas layer.register_parameter( "gemm1_clamp_limit", @@ -155,7 +179,11 @@ class TrtLlmNvFp4ExpertsBase: # raw. Register both on the layer so EPLB rearranges them with the # other per-expert tensors. if self.gemm1_beta is not None: - gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas + gemm1_beta = ( + self.gemm1_beta + if self.is_situ + else self.gemm1_beta / self.quant_config.g1_alphas + ) layer.register_parameter( "gemm1_beta", torch.nn.Parameter(gemm1_beta, requires_grad=False), @@ -197,13 +225,15 @@ class TrtLlmNvFp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI.""" + """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the + private SituGLU (SITU) cubin.""" return activation in [ MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.SITU, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 9622f9ed886..6339ab171e4 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -120,6 +120,8 @@ def FusedMoE( swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, e_score_correction_bias: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, activation: str = "silu", @@ -178,6 +180,8 @@ def FusedMoE( scoring_func: Scoring function for routing ("softmax" or others) routed_scaling_factor: Scaling factor applied to topk_weights or output swiglu_limit: SwiGLU activation limit + activation_situ_beta: SituGLU activation beta + activation_situ_linear_beta: SituGLU linear beta e_score_correction_bias: Expert score correction bias tensor apply_router_weight_on_input: Whether to apply router weights on input activation: Activation function name ("silu", "gelu", etc.) @@ -352,6 +356,8 @@ def FusedMoE( swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, skip_final_all_reduce=skip_final_all_reduce, ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 2c8ba874f0c..f7dc5dc2c05 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -890,6 +890,8 @@ class FusedMoEExpertsModular(FusedMoEExperts): beta: float = 0.0, topk_ids: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, ) -> None: apply_moe_activation( activation, @@ -900,6 +902,16 @@ class FusedMoEExpertsModular(FusedMoEExperts): beta=beta, topk_ids=topk_ids, expert_map=expert_map, + activation_situ_beta=( + self.moe_config.activation_situ_beta + if activation_situ_beta is None + else activation_situ_beta + ), + activation_situ_linear_beta=( + self.moe_config.activation_situ_linear_beta + if activation_situ_linear_beta is None + else activation_situ_linear_beta + ), ) @abstractmethod @@ -1307,6 +1319,14 @@ class FusedMoEKernelModularImpl: activation, ) + use_output_alias = ( + output_alias is not None + and output_alias.shape == fused_out.shape + and output_alias.dtype == fused_out.dtype + and output_alias.device == fused_out.device + and output_alias.is_contiguous() + ) + # If caller's output buffer already matches fused_out shape/dtype, alias # to skip the redundant copy in TopKWeightAndReduceNoOP.apply downstream. # This eliminates ~94% of __amd_rocclr_copyBuffer events (Copy 2 of the @@ -1314,15 +1334,10 @@ class FusedMoEKernelModularImpl: if current_platform.is_rocm(): from vllm._aiter_ops import rocm_aiter_ops - if ( - rocm_aiter_ops.is_fused_moe_enabled() - and output_alias is not None - and output_alias.shape == fused_out.shape - and output_alias.dtype == fused_out.dtype - and output_alias.device == fused_out.device - and output_alias.is_contiguous() - ): + if use_output_alias and rocm_aiter_ops.is_fused_moe_enabled(): fused_out = output_alias + elif use_output_alias: + fused_out = output_alias self.fused_experts.apply( output=fused_out, diff --git a/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py new file mode 100644 index 00000000000..5f881692215 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.utils.torch_utils import aux_stream, current_stream + +from .moe_runner import MoERunner, _unpack + +logger = init_logger(__name__) + + +class LatentMoERunner(MoERunner): + """MoE runner for latent MoE with a replicated routed up-projection. + + Fused path (tp>1, un-reduced combine output, shared expert, no SP): + concatenates the un-reduced latent partial (dim d) and the un-reduced + shared partial (dim D) into one contiguous buffer, all-reduces once, then + splits. The latent half is normed and up-projected locally (replicated + up-proj -> full hidden), and the shared add folds into the GEMM epilogue + (``torch.addmm``). One collective total, no post-reduction communication. + + Native path: the replicated up-proj produces the full hidden dim on every + rank, so the base runner combines routed + shared correctly at any TP size + (using two collectives instead of the fused path's one). + """ + + def __init__( + self, + *args, + enable_k3_latent_moe_tail_fusion: bool = False, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.enable_k3_latent_moe_tail_fusion = enable_k3_latent_moe_tail_fusion + use_fused_path = self._use_fused_path() + if ( + self.enable_k3_latent_moe_tail_fusion + and use_fused_path + and self.moe_config.tp_size not in (8, 16) + ): + logger.warning_once( + "K3 latent-MoE tail fusion currently supports TP=8 and TP=16, " + "but TP=%d is configured. Falling back to the default path.", + self.moe_config.tp_size, + ) + self.enable_k3_latent_moe_tail_fusion = False + + if self.enable_k3_latent_moe_tail_fusion and use_fused_path: + vllm_config = get_current_vllm_config() + if vllm_config.parallel_config.use_ubatching: + raise ValueError( + "K3 latent-MoE tail fusion does not support DBO or ubatching." + ) + if vllm_config.model_config.enable_sleep_mode: + raise ValueError( + "K3 latent-MoE tail fusion does not support sleep mode." + ) + transform = self.routed_output_transform + assert transform is not None + norm = transform.norm + assert norm is not None + from vllm.models.kimi_k3.nvidia.ops.latent_moe_tail import ( + KimiK3LatentMoETailOp, + ) + + op = KimiK3LatentMoETailOp.initialize( + hidden_size=transform.up_proj.weight.shape[0], + latent_size=norm.weight.shape[0], + dtype=norm.weight.dtype, + device=norm.weight.device, + rms_eps=norm.variance_epsilon, + ) + self._k3_latent_moe_tail_op = op + + def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel. + + flashinfer requires a residual buffer even when there is no residual to + add. + """ + numel = hidden_states.numel() + buf = getattr(self, "_zero_residual", None) + if ( + buf is None + or buf.dtype != hidden_states.dtype + or buf.device != hidden_states.device + or buf.numel() < numel + ): + buf = torch.zeros( + numel, dtype=hidden_states.dtype, device=hidden_states.device + ) + self._zero_residual = buf + return buf[:numel].view_as(hidden_states) + + def _use_fused_path(self) -> bool: + # The fused path merges the latent and shared reductions into one + # all-reduce, so it needs actual TP parallelism, a shared expert (to + # concat), an un-reduced combine output, and no sequence parallelism. + return ( + self.moe_config.tp_size > 1 + and self._shared_experts is not None + and not self._fused_output_is_reduced + and not self.moe_config.is_sequence_parallel + ) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + if self._use_fused_path(): + return self._fused_forward( + hidden_states, router_logits, input_ids, shared_experts_input + ) + return super().forward( + hidden_states, router_logits, input_ids, shared_experts_input + ) + + def _fused_forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, + ) -> torch.Tensor: + # When the caller pre-applies the routed input transform outside the + # runner (e.g. to overlap it on a separate stream), it passes the + # already-transformed routed input as ``hidden_states`` and the original + # hidden states as ``shared_experts_input``; skip the transform then. + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + hidden_states, og_hidden_dim_pre_xform, og_hidden_dim_post_xform = ( + self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + ) + + result = self._forward_entry( + hidden_states, + router_logits, + shared_experts_input, + input_ids, + self._encode_layer_name(), + self.moe_config.hidden_dim_unpadded + if self._quant_method.has_unpadded_output + else 0, + ) + + shared_output, fused_output = _unpack(result) + assert shared_output is not None + + if og_hidden_dim_pre_xform is not None: + fused_output = fused_output[..., :og_hidden_dim_pre_xform] + + transform = self.routed_output_transform + assert transform is not None + + if self.enable_k3_latent_moe_tail_fusion: + op = self._k3_latent_moe_tail_op + if 0 < fused_output.shape[0] <= op.contract.max_num_tokens: + norm = transform.norm + assert norm is not None + result = op( + fused_output, + shared_output, + norm.weight, + transform.up_proj.weight, + ) + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, output_is_reduced=True + ) + return self._maybe_add_zero_expert_output(result) + + fused_latent = None + if transform.norm is not None: + fused_latent = self.allreduce_norm_latent_out(fused_output, transform.norm) + else: + fused_latent = tensor_model_parallel_all_reduce(fused_output) + + shared_expert_stream = ( + aux_stream() + if shared_output.size(0) <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + else None + ) + if shared_expert_stream is not None: + # overlap shared expert allreduce with latent up_proj + main = current_stream() + shared_output.record_stream(shared_expert_stream) + shared_expert_stream.wait_stream(main) + with torch.cuda.stream(shared_expert_stream): + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + main.wait_stream(shared_expert_stream) + else: + shared_output = tensor_model_parallel_all_reduce(shared_output) + result = torch.mm(fused_latent, transform.up_proj.weight.t()) + result.add_(shared_output) + + # Output is already fully reduced; this only strips padding. + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, output_is_reduced=True + ) + + return self._maybe_add_zero_expert_output(result) + + def allreduce_norm_latent_out( + self, + hidden_states: torch.Tensor, + norm: RMSNorm, + ) -> tuple[torch.Tensor, torch.Tensor]: + """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.""" + if self.moe_config.tp_size == 1: + return norm(hidden_states) + + if flashinfer_trtllm_fused_allreduce_norm is not None: + ok, max_token_num = _can_use_flashinfer( + hidden_states, self.moe_config.tp_size + ) + if ok: + norm_out = torch.empty_like(hidden_states) + # With norm_out provided, the kernel writes the new residual + # (all_reduce(hidden_states) + residual) into the hidden_states + # buffer and the normalized result into norm_out. + flashinfer_trtllm_fused_allreduce_norm( + allreduce_in=hidden_states, + residual=self._get_zero_residual(hidden_states), + rms_gamma=norm.weight, + rms_eps=norm.variance_epsilon, + world_size=self.moe_config.tp_size, + weight_bias=0.0, + launch_with_pdl=True, + fp32_acc=True, + max_token_num=max_token_num, + pattern_code=_AR_RESIDUAL_RMS_NORM, + norm_out=norm_out, + ) + return norm_out + + reduced = tensor_model_parallel_all_reduce(hidden_states) + return norm(reduced) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 7942957f327..cdf436bcc5c 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -348,7 +348,8 @@ class MoERunner(MoERunnerInterface): return self.routed_experts.quant_method def apply_routed_input_transform( - self, hidden_states: torch.Tensor + self, + hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply transform for routed experts (e.g., latent projection). @@ -416,6 +417,7 @@ class MoERunner(MoERunnerInterface): def _maybe_reduce_shared_expert_output( self, shared_output: torch.Tensor | None, + fused_output_is_reduced: bool | None = None, ) -> torch.Tensor | None: """All-reduce shared expert output when the combine kernel already reduced fused output. @@ -425,18 +427,43 @@ class MoERunner(MoERunnerInterface): * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled in the model. """ + if fused_output_is_reduced is None: + fused_output_is_reduced = self._fused_output_is_reduced + if ( shared_output is not None and not self.moe_config.is_sequence_parallel - and self._fused_output_is_reduced + and fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) return shared_output + def _maybe_reduce_routed_output_before_transform( + self, + fused_output: torch.Tensor, + fused_output_is_reduced: bool, + ) -> tuple[torch.Tensor, bool]: + """All-reduce latent routed output before its output transform. + + Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm. + TP partial routed outputs must be summed in latent space before such + transforms are applied. + """ + if ( + self.routed_output_transform is not None + and not self.moe_config.is_sequence_parallel + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + and not fused_output_is_reduced + ): + fused_output = tensor_model_parallel_all_reduce(fused_output) + fused_output_is_reduced = True + return fused_output, fused_output_is_reduced + def _maybe_reduce_final_output( self, states: torch.Tensor, trunc_size: int | None, + output_is_reduced: bool | None = None, ) -> torch.Tensor: """All-reduce the combined output if needed. @@ -455,11 +482,14 @@ class MoERunner(MoERunnerInterface): # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. + if output_is_reduced is None: + output_is_reduced = self._fused_output_is_reduced + if ( not self.moe_config.is_sequence_parallel and not self.moe_config.skip_final_all_reduce and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - and not self._fused_output_is_reduced + and not output_is_reduced ): states = tensor_model_parallel_all_reduce(states) @@ -643,6 +673,7 @@ class MoERunner(MoERunnerInterface): hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: """Invoke the fused moe layer. @@ -664,11 +695,16 @@ class MoERunner(MoERunnerInterface): _moe_forward and _moe_forward_shared must be split. """ - # Apply transform for routed experts (e.g., latent projection - # for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) + # Apply transform for routed experts (e.g., latent projection for + # latent MoE). When the caller pre-applies the routed input transform + # outside the runner (e.g. to overlap it on a separate stream), it + # passes the already-transformed routed input as ``hidden_states`` and + # the original hidden states as ``shared_experts_input``; skip the + # transform in that case so shared experts still see the original input. + if shared_experts_input is None: + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) # Record before `_maybe_pad_hidden_states` pads activations to match # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi` @@ -708,9 +744,22 @@ class MoERunner(MoERunnerInterface): if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] - # If combine kernel already reduced fused, reduce shared to match. + fused_output_is_reduced = self._fused_output_is_reduced + + # Latent routed output has to be reduced before output transform, + # because the transform may include non-linear normalization. + fused_output, fused_output_is_reduced = ( + self._maybe_reduce_routed_output_before_transform( + fused_output, + fused_output_is_reduced, + ) + ) + + # If routed output is already reduced, reduce shared to match. # See note above re: the two all-reduce points. - shared_output = self._maybe_reduce_shared_expert_output(shared_output) + shared_output = self._maybe_reduce_shared_expert_output( + shared_output, fused_output_is_reduced + ) shared_output, fused_output = self._maybe_apply_routed_scale_to_output( shared_output, fused_output @@ -724,7 +773,9 @@ class MoERunner(MoERunnerInterface): else: result = fused_output - result = self._maybe_reduce_final_output(result, og_hidden_dim_post_xform) + result = self._maybe_reduce_final_output( + result, og_hidden_dim_post_xform, fused_output_is_reduced + ) return self._maybe_add_zero_expert_output(result) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py index cc79095ead8..5b8e0a3bff3 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_interface.py @@ -36,6 +36,7 @@ class MoERunnerInterface(PluggableLayer, ABC): hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: torch.Tensor | None = None, ) -> torch.Tensor: raise NotImplementedError diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 15f08f26550..ddeffb56457 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1080,6 +1080,7 @@ def causal_conv1d_update( block_idx_last_scheduled_token: torch.Tensor | None = None, initial_state_idx: torch.Tensor | None = None, validate_data=False, + out: torch.Tensor | None = None, ): """ x: Input tensor which can take the following shapes: @@ -1117,7 +1118,8 @@ def causal_conv1d_update( for example: conv_state_indices = [null_block_id, 1, 20, null_block_id] in this case, the kernel will not process entries at indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x` + out: optional output tensor with the same shape as `x`. When omitted, + the input is overwritten. """ if validate_data: assert null_block_id is not None @@ -1129,10 +1131,22 @@ def causal_conv1d_update( original_x_dtype = x.dtype x = x.to(conv_state.dtype) + if out is None: + out = x + else: + if out.shape != x.shape: + raise ValueError( + f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}." + ) + if out.dtype != original_x_dtype or out.device != x.device: + raise ValueError( + "`out` must have the same dtype and device as the input `x`." + ) unsqueeze = query_start_loc is None and x.dim() == 2 if unsqueeze: # make it (batch, dim, seqlen) with seqlen == 1 x = x.unsqueeze(-1) + out = out.unsqueeze(-1) if query_start_loc is None: batch, dim, seqlen = x.shape else: @@ -1159,8 +1173,6 @@ def causal_conv1d_update( assert num_cache_lines >= batch assert weight.stride(1) == 1 # Need this - # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o' - out = x stride_w_dim, stride_w_width = weight.stride() if query_start_loc is None: diff --git a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py new file mode 100644 index 00000000000..b952e3ebbce --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _gather_initial_states_kernel( + state_ptr, + indices_ptr, + has_initial_state_ptr, + output_ptr, + stride_state_batch, + stride_indices, + stride_has_initial_state, + row_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + block_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < row_size + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + has_initial_state = tl.load( + has_initial_state_ptr + batch_idx * stride_has_initial_state + ).to(tl.int1) + state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) + state_idx = tl.where(has_initial_state, state_idx, 0) + values = tl.load( + state_ptr + state_idx * stride_state_batch + offsets, + mask=mask & has_initial_state, + other=0.0, + ) + tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask) + + +def gather_initial_states( + state: torch.Tensor, + indices: torch.Tensor, + has_initial_state: torch.Tensor, +) -> torch.Tensor: + """Gather dense state rows, replacing uninitialized rows with zeros.""" + assert state.ndim >= 2 + assert state.is_cuda + assert indices.ndim == 1 and has_initial_state.ndim == 1 + assert indices.shape == has_initial_state.shape + assert indices.device == state.device + assert has_initial_state.device == state.device + assert indices.dtype in (torch.int32, torch.int64) + assert has_initial_state.dtype == torch.bool + + row_size = state[0].numel() + # Mamba pages may pad stride(0), but each state row remains dense. + assert state[0].is_contiguous() + output = torch.empty( + (indices.numel(), *state.shape[1:]), + dtype=state.dtype, + device=state.device, + ) + block_size = min(triton.next_power_of_2(row_size), 1024) + grid = (triton.cdiv(row_size, block_size), indices.numel()) + _gather_initial_states_kernel[grid]( + state, + indices, + has_initial_state, + output, + state.stride(0), + indices.stride(0), + has_initial_state.stride(0), + row_size=row_size, + BLOCK_SIZE=block_size, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output diff --git a/vllm/models/common/__init__.py b/vllm/models/common/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/common/ops/__init__.py b/vllm/models/common/ops/__init__.py new file mode 100644 index 00000000000..0d0fe8de26e --- /dev/null +++ b/vllm/models/common/ops/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ops shared across model implementations.""" + +from .fused_qk_rmsnorm import fused_q_kv_rmsnorm + +__all__ = [ + "fused_q_kv_rmsnorm", +] diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py new file mode 100644 index 00000000000..ac7d0dfd4a1 --- /dev/null +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_q_kv_rmsnorm_kernel( + q_ptr, + q_out_ptr, + q_weight_ptr, + q_in_stride, + q_out_stride, + kv_ptr, + kv_out_ptr, + kv_weight_ptr, + kv_in_stride, + kv_out_stride, + eps, + Q_SIZE: tl.constexpr, + KV_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +): + # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y. + # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes + # the launch at max-num-batched-tokens >= 65536 with "invalid argument". + # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32 + # past num_tokens ~87K under large chunked prefill. + token_idx = tl.program_id(0).to(tl.int64) + pid_task = tl.program_id(1) + + if pid_task == 0: + SIZE = Q_SIZE + row_in = q_ptr + token_idx * q_in_stride + weight_ptr = q_weight_ptr + row_out = q_out_ptr + token_idx * q_out_stride + else: + SIZE = KV_SIZE + row_in = kv_ptr + token_idx * kv_in_stride + weight_ptr = kv_weight_ptr + row_out = kv_out_ptr + token_idx * kv_out_stride + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's + # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which + # keep x, rrms, and w all in fp32 and perform a single cast at store. + block = tl.arange(0, BLOCK_SIZE) + mask = block < SIZE + x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / SIZE + rrms = tl.rsqrt(variance + eps) + w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32) + y = x * rrms * w + tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask) + + +def fused_q_kv_rmsnorm( + qr: torch.Tensor, + kv: torch.Tensor, + q_weight: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + assert qr.ndim == 2 and kv.ndim == 2 + assert qr.shape[0] == kv.shape[0], ( + f"token dim mismatch: qr={qr.shape}, kv={kv.shape}" + ) + assert qr.stride(-1) == 1 and kv.stride(-1) == 1 + assert q_weight.is_contiguous() and kv_weight.is_contiguous() + + q_size = qr.shape[1] + kv_size = kv.shape[1] + num_tokens = qr.shape[0] + qr_out = torch.empty_like(qr) + kv_out = torch.empty_like(kv) + if num_tokens == 0: + return qr_out, kv_out + + block_size = triton.next_power_of_2(max(q_size, kv_size)) + _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)]( + qr, + qr_out, + q_weight, + qr.stride(0), + qr_out.stride(0), + kv, + kv_out, + kv_weight, + kv.stride(0), + kv_out.stride(0), + eps, + Q_SIZE=q_size, + KV_SIZE=kv_size, + BLOCK_SIZE=block_size, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return qr_out, kv_out diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py index f547f650629..883545e3f2d 100644 --- a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -130,13 +130,20 @@ def inkling_fa4_rel_attention( cute_window = (None, None) if window_size == (-1, -1) else window_size rel_logits = rel_logits.contiguous() + flash_attn_varlen_func: Callable[..., Any] if _use_sheared_bias(): - from vllm.third_party.tml_fa4 import flash_attn_varlen_func + from vllm.third_party.tml_fa4 import ( + flash_attn_varlen_func as tml_flash_attn_varlen_func, + ) + flash_attn_varlen_func = tml_flash_attn_varlen_func bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} else: - from vllm.vllm_flash_attn.cute import flash_attn_varlen_func + from vllm.vllm_flash_attn.cute import ( + flash_attn_varlen_func as cute_flash_attn_varlen_func, + ) + flash_attn_varlen_func = cute_flash_attn_varlen_func bias_kwargs = { "score_mod": _get_score_mod(rel_extent), "aux_tensors": [rel_logits], diff --git a/vllm/models/kimi_k3/__init__.py b/vllm/models/kimi_k3/__init__.py index 208f01a7cb5..54f965eadaa 100644 --- a/vllm/models/kimi_k3/__init__.py +++ b/vllm/models/kimi_k3/__init__.py @@ -1,2 +1,28 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi K3 model — hardware-isolated entry point. + +The implementation lives under ``nvidia/`` and ``amd/``; this module picks the +right one for the current platform and re-exports the public classes used by +the model registry. (Mirrors ``vllm.models.minimax_m3``.) +""" + +from typing import TYPE_CHECKING + +from vllm.platforms import current_platform + +# The NVIDIA branch is the static default that type-checkers see; the ROCm +# branch overrides it at runtime (kept type-compatible via type: ignore). +if TYPE_CHECKING or not current_platform.is_rocm(): + from .nvidia.model import KimiK3ForConditionalGeneration, KimiLinearForCausalLM + from .nvidia.mtp import KimiK3MTP +else: + from .amd.linear import KimiLinearForCausalLM # type: ignore[assignment] + from .amd.model import KimiK3ForConditionalGeneration # type: ignore[assignment] + from .amd.mtp import KimiK3MTP # type: ignore[assignment] + +__all__ = [ + "KimiK3ForConditionalGeneration", + "KimiK3MTP", + "KimiLinearForCausalLM", +] diff --git a/vllm/models/kimi_k3/amd/linear.py b/vllm/models/kimi_k3/amd/linear.py new file mode 100644 index 00000000000..f27c81a24f1 --- /dev/null +++ b/vllm/models/kimi_k3/amd/linear.py @@ -0,0 +1,1065 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + get_spec_layer_idx_from_weight_name, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.models.kimi_k3.amd.ops.attn_res import attn_res +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.math_utils import cdiv + +logger = init_logger(__name__) + + +class KimiMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu and situ are supported." + ) + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +def _apply_attn_res( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: ReplicatedLinear, + norm: RMSNorm, + num_valid_blocks: int, +) -> torch.Tensor: + if num_valid_blocks <= 0: + return prefix_sum + + return attn_res( + prefix_sum, + block_residual, + norm.weight, + proj.weight.squeeze(0), + num_valid_blocks, + norm.variance_epsilon, + ) + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = ( + routed_expert_hidden_size + if routed_expert_hidden_size is not None + else hidden_size + ) + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr( + config, "min_moe_intermediate_per_partition", 256 + ) + if self.tp_size > 1: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = ( + min_moe_intermediate_per_partition * self.tp_size + ) + activation_situ_beta = ( + config.activation_situ_beta if config.hidden_act == "situ" else None + ) + activation_situ_linear_beta = ( + config.activation_situ_linear_beta if config.hidden_act == "situ" else None + ) + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) + + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if self.latent_moe_use_norm + else None + ) + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + routed_input_transform=self.routed_expert_down_proj, + routed_output_transform=self.routed_output_transform, + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = self.experts.w13_weight_packed + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = self.experts.w2_weight_packed + w13_weight.data.zero_() + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = ( + moe_intermediate_size // self.tp_size + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiMLAAttention(nn.Module): + """ + Main reference: DeepseekV2 vllm Implementation + """ + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_nope: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + **kwargs, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + self.num_local_heads = num_heads // tp_size + self.scaling = self.qk_head_dim**-0.5 + self.use_nope = use_nope + assert self.use_nope is True + assert num_heads % tp_size == 0 + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm( + self.q_lora_rank, + eps=config.rms_norm_eps, + ) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm( + self.kv_lora_rank, + eps=config.rms_norm_eps, + ) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.use_output_gate = config.mla_use_output_gate + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = ColumnParallelLinear( + self.hidden_size, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + mla_modules = MLAModules( # type: ignore[call-arg] + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=None, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=None, + is_sparse=False, + topk_indices_buffer=None, + g_proj=getattr(self, "g_proj", None), + ) + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + output: torch.Tensor, + ) -> None: + output[:] = self.mla_attn(positions, hidden_states) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + if config.is_kda_layer(layer_idx): + self.self_attn = KimiGatedDeltaNetAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope is not None + self.self_attn = KimiMLAAttention( + layer_idx=layer_idx, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + quant_config=quant_config, + cache_config=cache_config, + model_config=model_config, + prefix=f"{prefix}.self_attn", + config=config, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_nope=mla_use_nope, + ) + + if ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ): + self.block_sparse_moe = KimiMoE( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_residuals = attn_res_block_size is not None + if attn_res_block_size is not None: + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + attn_output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=attn_output, + ) + return attn_output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.use_attn_residuals: + assert residual is not None + return self.forward_attn_residual(positions, hidden_states, residual) + + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self._run_self_attn(positions, hidden_states) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + def forward_attn_residual( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + prefix_sum = hidden_states + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.prev_valid_blocks, + ) + + if self.is_block_write_layer: + block_residual[:, self.block_write_idx, :].copy_(prefix_sum) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self._run_self_attn(positions, hidden_states) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + mlp_valid_blocks = self.prev_valid_blocks + ( + 1 if self.is_block_write_layer else 0 + ) + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + mlp_valid_blocks, + ) + + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + prefix_sum = prefix_sum + hidden_states + return prefix_sum, block_residual + + +class KimiLinearModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + self.config = config + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + vllm_config, + prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.attn_res_block_size is not None: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if config.attn_res_block_size is not None: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.config.attn_res_block_size is not None: + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.config.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def _maybe_add_hidden_state( + self, + aux_hidden_states: list[torch.Tensor], + layer_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> list[torch.Tensor]: + if self.config.attn_res_block_size is not None: + # attn-res `residual` is a block-state bank, not an additive + # residual; None makes the mixin capture the prefix sum directly. + residual = None + return super()._maybe_add_hidden_state( + aux_hidden_states, layer_idx, hidden_states, residual + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, residual + ) + + if self.config.attn_res_block_size is None: + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + # NOTE: the final norm is applied in compute_logits instead of here, + # so the MTP draft model receives the pre-norm hidden states. + if residual is not None: + hidden_states = hidden_states + residual + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + attn_res_block_num = cdiv(self.end_layer, self.config.attn_res_block_size) + block_residual = hidden_states.new_empty( + hidden_states.size(0), attn_res_block_num, hidden_states.size(1) + ) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + # AMD attn-res layer already returns prefix_sum + MLP delta as + # hidden_states; the override drops the block bank in residual. + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states = _apply_attn_res( + hidden_states, + residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + attn_res_block_num, + ) + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[ + tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]] + ], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + if self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + +class KimiLinearForCausalLM( + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + return hidden_states + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + return loader.load_weights(weights) diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py new file mode 100644 index 00000000000..ec49891ddb2 --- /dev/null +++ b/vllm/models/kimi_k3/amd/model.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +from collections.abc import Iterable +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from vllm.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) +from .linear import KimiLinearForCausalLM + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + quant_config = vllm_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel( + config.vision_config.num_attention_heads + ) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(vllm_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to( + device=self.device, dtype=model_config.dtype + ) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to( + device=self.device, dtype=model_config.dtype + ) + + self.quant_config = quant_config + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + def _maybe_ignore_quant_config( + self, quant_config: QuantizationConfig | None + ) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input( + self, **kwargs: object + ) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError( + "pixel_values must be a tensor or a list of tensors, " + f"got {type(pixel_values)}" + ) + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] + ) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), ( + f"expect grid_thws to be a tensor, got {type(grid_thws)}" + ) + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( + f"unexpected shape for grid_thws: {grid_thws.shape}" + ) + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input( + self, media_input: KimiK25MediaPixelInputs + ) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + hidden_states = self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs + ) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( + batch_size + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/kimi_k3/amd/mtp.py b/vllm/models/kimi_k3/amd/mtp.py new file mode 100644 index 00000000000..8fd79d5a149 --- /dev/null +++ b/vllm/models/kimi_k3/amd/mtp.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .linear import KimiDecoderLayer, get_spec_layer_idx_from_weight_name + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + # The MTP block is a standard KimiDecoderLayer, but it must NOT use the + # attn-residual (block-residual) scheme even when the base model does: + # the draft starts from the target's hidden state, without the main + # model's accumulated per-block residual tensor. We disable it by + # shallow-copying the config with ``attn_res_block_size=None``. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + self.mtp_block = KimiDecoderLayer(block_config, vllm_config, prefix=prefix) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + + hidden_states, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm( + hidden_states, residual + ) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer( + config, vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + expert_params_mapping = ( + fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + if self.config.is_moe + else [] + ) + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not any(n.endswith("w13_weight_packed") for n in params_dict) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and ( + ".layers" not in name + ): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/vllm/models/kimi_k3/amd/ops/third_party/__init__.py b/vllm/models/kimi_k3/amd/ops/third_party/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py new file mode 100644 index 00000000000..77ed09d489e --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# AMD/ROCm vendored copy of the Kimi-K3 KDA triton kernels. +# +# Provenance: mirror of vllm/models/kimi_k3/nvidia/ops/third_party/kda (tracker +# mke-tracker @ 7adebfcf9; FLA vendored per PRs #39/#86). Split per-vendor so +# AMD can carry gfx950-specific kernel changes without touching the NVIDIA copy. +# +# fla-org/flash-linear-attention#869 (the unmerged ROCm fixes our earlier +# amd_fla shim carried against FLA 0.5.0) is covered here by the *newer* vendored +# FLA rather than the literal patch: +# - transpose-state-layout workaround: N/A (kernels rewritten; no +# transpose_state_layout path remains), +# - AMD autotune configs: present (is_amd num_warps/num_stages branches), +# - OOB-mask correctness fix: present (all tl.load use mask=..., other=0). +# Validated on gfx950: no core-dump, gsm8k 94.1%. +# +# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync +# with the NVIDIA copy on FLA updates; any divergence should be an intentional, +# documented gfx950-specific change (a #869-style AMD-only fix). + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py new file mode 100644 index 00000000000..7d0ea0204ce --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py @@ -0,0 +1,935 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.op import exp2, log +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError(f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}") + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py new file mode 100644 index 00000000000..ca3bc8f75e3 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,662 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.op import exp2, gather +from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported +from vllm.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_g0 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_k1 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_g1 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_k2 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_g2 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_k3 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + p_g3 = tl.make_block_ptr( + g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to( + tl.float32 + ) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b1 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,) + ) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Aqk21 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b2 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,) + ) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Aqk31 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Aqk32 = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_b3 = tl.make_block_ptr( + beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,) + ) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.0) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2 * BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.0) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.0) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) + for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV * BC + o_i) + b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.0) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk10 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Akk20 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk21 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk30 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Akk31 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Akk32 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BK", "NC", "BT", "HV"], +) +@triton.jit(do_not_specialize=["B", "T"]) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather( + b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0 + ) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV * K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.0) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.0) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr( + Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0) + ) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV * BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.0) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 00000000000..cc448d95c0e --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp2 +from vllm.triton_utils import tl, triton + + +@triton.heuristics( + { + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BH": BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=["T", "N"]) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H * K + k += bos * H * K + g += bos * HV * K + Aqk += bos * HV * BT + Akk += bos * HV * BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr( + g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr( + g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0) + ) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store( + Aqk + i_t * HV * BT + o_hv * BT + j % BT, + b_Aqk.to(Aqk.dtype.element_ty), + mask=m_hv, + ) + tl.store( + Akk + i_t * HV * BC + o_hv * BC + j - i_ts, + b_Akk.to(Akk.dtype.element_ty), + mask=m_hv, + ) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): + return (B * T, triton.cdiv(HV, meta["BH"])) + + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 00000000000..2f512df6264 --- /dev/null +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,621 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp, log +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, +): + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + BV = 32 if use_gate_in_kernel else 8 + num_warps = 4 if use_gate_in_kernel else 1 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=2, + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + # gfx950: always fuse the gate. + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + stride_state_token: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + state_idx = tl.load(state_indices + i_n).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1 or state_indices.stride(0) != 1: + raise ValueError("`state_indices` must be contiguous and one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + ) + return out, initial_state diff --git a/vllm/models/kimi_k3/common/__init__.py b/vllm/models/kimi_k3/common/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/common/mm_preprocess.py b/vllm/models/kimi_k3/common/mm_preprocess.py new file mode 100644 index 00000000000..fa70ccdce0a --- /dev/null +++ b/vllm/models/kimi_k3/common/mm_preprocess.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared Kimi-K3 multimodal preprocessing.""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import torch +from transformers import BatchFeature + +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + InputProcessingContext, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config +from vllm.transformers_utils.processor import cached_get_image_processor +from vllm.transformers_utils.processors.kimi_k3 import KimiK3Processor + +logger = init_logger(__name__) + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +): + # Apply the patch limits. + s1 = math.sqrt( + in_patch_limit + / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) + ) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, patch_limit_on_one_side * patch_size) + + factor = merge_kernel_size * patch_size + + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > " + f"patch_limit_on_one_side {patch_limit_on_one_side}" + ) + + num_tokens = token_height * token_width + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +class KimiK3ProcessingInfo(BaseProcessingInfo): + """Processing information for the image-only Kimi-K3 model. + + K3 uses the standard ``image`` modality (unlike K2.5's unified + ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around + the checkpoint's image processor and resolves the ``<|media_pad|>`` token + id the same way K2.5 does. + """ + + def __init__(self, ctx: InputProcessingContext) -> None: + super().__init__(ctx) + + self.hf_config = hf_config = self.get_hf_config() + + tokenizer = self.get_tokenizer() + image_processor = cached_get_image_processor( + self.ctx.model_config.model, + revision=self.ctx.model_config.revision, + trust_remote_code=self.ctx.model_config.trust_remote_code, + ) + + # Resolve token ID from the tokenizer because transformers v5 + # may remap token IDs vs config.json. + config_token_id = hf_config.media_placeholder_token_id + resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") + unk_token_id = getattr(tokenizer, "unk_token_id", None) + is_valid_resolved = isinstance(resolved_token_id, int) and ( + unk_token_id is None or resolved_token_id != unk_token_id + ) + if is_valid_resolved and resolved_token_id != config_token_id: + logger.warning_once( + "Kimi-K3 config.media_placeholder_token_id (%d) disagrees " + "with tokenizer mapping for <|media_pad|> (%d). " + "Using tokenizer value.", + config_token_id, + resolved_token_id, + ) + media_token_id = resolved_token_id + # Patch config so downstream code also sees the correct ID. + hf_config.media_placeholder_token_id = resolved_token_id + else: + media_token_id = config_token_id + + self.media_token_id = media_token_id + self.media_token = tokenizer.decode(media_token_id) + + self.image_processor = image_processor + self.hf_processor = KimiK3Processor( + tokenizer=tokenizer, + image_processor=image_processor, + ) + self.media_tokens_calculator = image_processor.media_tokens_calculator + + def get_hf_processor(self, **kwargs: object) -> KimiK3Processor: + return self.hf_processor + + def get_hf_config(self) -> KimiK3Config: + return self.ctx.get_hf_config(KimiK3Config) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # None means unlimited + return {"image": None} + + @classmethod + def get_max_image_size( + cls, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, + ) -> ImageSize: + max_side = patch_limit_on_one_side * patch_size + best_score = (-1, -1) + best_size = (max_side, max_side) + + for width_patches in range(patch_limit_on_one_side + 1): + width = min((width_patches + 1) * patch_size - 1, max_side) + for height_patches in range(width_patches, patch_limit_on_one_side + 1): + height = min((height_patches + 1) * patch_size - 1, max_side) + resize_config = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit, + patch_limit_on_one_side, + fixed_output_tokens, + ) + padded_width = resize_config["new_width"] + resize_config["pad_width"] + padded_height = ( + resize_config["new_height"] + resize_config["pad_height"] + ) + num_patches = padded_width // patch_size * (padded_height // patch_size) + score = (resize_config["num_tokens"], num_patches) + if score > best_score: + best_score = score + best_size = (width, height) + return ImageSize(width=best_size[0], height=best_size[1]) + + +class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]): + """Builds image-based dummy inputs for K3 profiling. + + The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly + the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy + mm data is a plain list of PIL images under the ``image`` key. + """ + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + return self.info.get_hf_config().image_placeholder * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + media_proc_cfg = self.info.image_processor.media_proc_cfg + max_size = self.info.get_max_image_size( + media_proc_cfg["patch_size"], + media_proc_cfg["merge_kernel_size"], + media_proc_cfg["in_patch_limit"], + media_proc_cfg["patch_limit_on_one_side"], + media_proc_cfg["fixed_output_tokens"], + ) + num_images = mm_counts.get("image", 0) + image_overrides = cast( + ImageDummyOptions | None, + mm_options.get("image") if mm_options else None, + ) + return { + "image": self._get_dummy_images( + width=max_size.width, + height=max_size.height, + num_images=num_images, + overrides=image_overrides, + ) + } + + +class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]): + """Image-only multi-modal processor for Kimi-K3.""" + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Override so the base always routes through the text+mm path + # (`KimiK3Processor.__call__`). Otherwise the mm-only fast path calls + # the checkpoint image processor directly with bare PIL images, but it + # requires `{"type": "image", "image": PIL}` media dicts that only our + # wrapper builds. + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + """Slice the flattened patch tensor back into per-image items. + + ``pixel_values`` holds all patches from every image concatenated; each + image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one + ``[N_t, N_h, N_w]`` row per image. + """ + grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3))) + grid_sizes = grid_thws.prod(-1) + + return dict( + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes), + grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + """Expand each K3 image placeholder into a resolution-aware update. + + K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per + image. This replaces that token with + ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``, + embedding the per-image resolution in the prompt and marking only the + ``<|media_pad|>`` positions as embedding slots (the number of pads is + the feature size returned by ``media_tokens_calculator``). + """ + media_token_id = self.info.media_token_id + media_token = self.info.media_token + image_placeholder = self.info.get_hf_config().image_placeholder + + def get_replacement(item_idx: int) -> PromptUpdateDetails[str]: + images = mm_items.get_items("image", ImageProcessorItems) + image = images.get(item_idx) + if image is None: + raise ValueError(f"Missing image data at index {item_idx}") + + # The checkpoint image processor works on media dicts, so wrap the + # PIL image before asking it for the token count. + num_media_token = self.info.media_tokens_calculator( + {"type": "image", "image": image} + ) + pads = media_token * num_media_token + + # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the + # post-preprocess (smart-resized) ones. `image` comes from the + # untouched parsed `mm_items`; the checkpoint image processor + # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via + # `image.resize(...)` and never mutates the stored PIL. This matches + # the reference HF processor (`KimiK3Processor.preprocess_medias`), + # which also builds the prompt from the original `img.size`. The + # resize is reflected only in the pad count above. + width, height = images.get_image_size(item_idx) + full = ( + f"<|media_begin|>image {width}x{height}<|media_content|>" + f"{pads}<|media_end|>" + ) + + return PromptUpdateDetails.select_token_id(full, media_token_id) + + return [ + PromptReplacement( + modality="image", + target=image_placeholder, + replacement=get_replacement, + ), + ] diff --git a/vllm/models/kimi_k3/common/mtp.py b/vllm/models/kimi_k3/common/mtp.py new file mode 100644 index 00000000000..76c796c1b0b --- /dev/null +++ b/vllm/models/kimi_k3/common/mtp.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused Kimi-K3 MTP input preparation.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _rms_norm(x, weight, eps, hidden_size: tl.constexpr): + x = x.to(tl.float32) + variance = tl.sum(x * x, axis=0) / hidden_size + return x * tl.rsqrt(variance + eps) * weight.to(tl.float32) + + +@triton.jit +def _fused_mtp_input_kernel( + positions_ptr, + inputs_embeds_ptr, + previous_hidden_states_ptr, + enorm_weight_ptr, + hnorm_weight_ptr, + output_ptr, + eps, + inputs_embeds_stride, + previous_hidden_states_stride, + output_stride, + hidden_size: tl.constexpr, + block_size: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + input_idx = tl.program_id(1) + offsets = tl.arange(0, block_size) + mask = offsets < hidden_size + + if input_idx == 0: + position = tl.load(positions_ptr + token_idx) + values = tl.load( + inputs_embeds_ptr + token_idx * inputs_embeds_stride + offsets, + mask=mask & (position != 0), + other=0.0, + ) + weight = tl.load(enorm_weight_ptr + offsets, mask=mask, other=0.0) + else: + values = tl.load( + previous_hidden_states_ptr + + token_idx * previous_hidden_states_stride + + offsets, + mask=mask, + other=0.0, + ) + weight = tl.load(hnorm_weight_ptr + offsets, mask=mask, other=0.0) + + output = _rms_norm(values, weight, eps, hidden_size) + tl.store( + output_ptr + token_idx * output_stride + input_idx * hidden_size + offsets, + output, + mask=mask, + ) + + +def fused_mtp_input( + positions: torch.Tensor, + inputs_embeds: torch.Tensor, + previous_hidden_states: torch.Tensor, + enorm_weight: torch.Tensor, + hnorm_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Mask and normalize both MTP inputs into the projection layout.""" + num_tokens, hidden_size = inputs_embeds.shape + output = torch.empty( + num_tokens, + 2 * hidden_size, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + if num_tokens == 0: + return output + + _fused_mtp_input_kernel[(num_tokens, 2)]( + positions, + inputs_embeds, + previous_hidden_states, + enorm_weight, + hnorm_weight, + output, + eps, + inputs_embeds.stride(0), + previous_hidden_states.stride(0), + output.stride(0), + hidden_size, + triton.next_power_of_2(hidden_size), + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py new file mode 100644 index 00000000000..f9692d5d472 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3 dense MLA draft model for DSpark speculative decoding.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import vllm._custom_ops as ops +from vllm.config import VllmConfig +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + get_draft_quant_config, + maybe_prefix, +) +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.model import KimiMLP +from vllm.utils.torch_utils import is_quantized_kv_cache + + +class ReplicatedDSparkMarkovHead(DSparkMarkovHead): + """DSpark Markov head with full weights on every TP rank.""" + + def __init__( + self, + vocab_size: int, + draft_vocab_size: int, + markov_rank: int, + prefix: str, + ) -> None: + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + super().__init__( # type: ignore[call-arg] + vocab_size, + draft_vocab_size, + markov_rank, + prefix, + replicated=True, + ) + + +class K3DSparkDecoderLayer(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + config, + layer_idx: int, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + quant_config = get_draft_quant_config(vllm_config) + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=maybe_prefix( + prefix, f"layers.{start_layer_id + layer_idx}.self_attn" + ), + use_rope=True, + non_causal_multi_token_decode=True, + ) + self.mlp = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class K3DSparkModel(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.quant_config = get_draft_quant_config(vllm_config) + + # The frozen target embedding is aliased after the draft checkpoint loads. + self.embed_tokens: nn.Module | None = None + + self.context_proj = ReplicatedLinear( + self.config.target_hidden_size * self.config.num_target_layers, + self.config.hidden_size, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "context_proj"), + ) + self.context_norm = RMSNorm( + self.config.hidden_size, eps=self.config.rms_norm_eps + ) + + self.layers = nn.ModuleList( + [ + K3DSparkDecoderLayer( + vllm_config=vllm_config, + config=self.config, + layer_idx=layer_idx, + start_layer_id=start_layer_id, + prefix=prefix, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.markov_head = ReplicatedDSparkMarkovHead( + self.config.vocab_size, + self.config.draft_vocab_size, + self.config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + self._context_kv_fusion_available: bool | None = None + self._max_num_context_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + assert self.embed_tokens is not None + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.context_norm(self.context_proj(hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project target-derived context into each draft layer's latent cache.""" + if self._context_kv_fusion_available is None: + self._build_fused_context_kv_buffers() + if self._context_kv_fusion_available: + self._precompute_fused_context_kv( + context_states, context_positions, context_slot_mapping + ) + return + + # Quantized fallback. Directly invoking the projection modules preserves + # their quantization methods, at the cost of also computing unused Q rows. + for layer_idx, layer in enumerate(self.layers): + attn = layer.self_attn + assert attn.fused_qkv_a_proj is not None + assert attn.q_lora_rank is not None + assert attn.rotary_emb is not None + qkv_lora = attn.fused_qkv_a_proj(context_states)[0] + kv_lora = qkv_lora[..., attn.q_lora_rank :] + kv_c, k_pe = kv_lora.split( + [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1 + ) + kv_c = attn.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors. + # The vLLM CUDA op supports rotating one tensor in place and + # consumes the same (possibly scaled fp32) cos/sin cache. + rotary_emb = attn.rotary_emb + ops.rotary_embedding( + context_positions, + k_pe, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn.impl.do_kv_cache_update( + kv_c, + k_pe, + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _build_fused_context_kv_buffers(self) -> None: + """Build a cross-layer KV-only A projection after checkpoint loading.""" + if self.quant_config is not None: + self._context_kv_fusion_available = False + return + + attentions = [layer.self_attn for layer in self.layers] + if not attentions or any( + attn.fused_qkv_a_proj is None + or not hasattr(attn.fused_qkv_a_proj, "weight") + for attn in attentions + ): + self._context_kv_fusion_available = False + return + + attn0 = attentions[0] + assert attn0.q_lora_rank is not None + kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim + kv_weights = [] + for attn in attentions: + assert attn.q_lora_rank is not None + assert ( + attn.q_lora_rank == attn0.q_lora_rank + and attn.kv_lora_rank == attn0.kv_lora_rank + and attn.qk_rope_head_dim == attn0.qk_rope_head_dim + and attn.kv_a_layernorm.variance_epsilon + == attn0.kv_a_layernorm.variance_epsilon + ), "All MLA DSpark layers must share their latent KV geometry." + kv_weights.append( + attn.fused_qkv_a_proj.weight.detach().narrow( + 0, attn.q_lora_rank, kv_width + ) + ) + + # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused + # A weights are replicated (`disable_tp=True`), so this is valid on + # every TP rank without communication. + self._fused_context_kv_weight = torch.cat(kv_weights, dim=0) + self._context_kv_norm_weights = torch.stack( + [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 + ).contiguous() + self._num_context_layers = len(attentions) + self._context_kv_width = kv_width + self._context_kv_lora_rank = attn0.kv_lora_rank + self._context_rope_dim = attn0.qk_rope_head_dim + self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._context_positions_repeated = torch.empty( + self._num_context_layers * self._max_num_context_tokens, + dtype=torch.int64, + device=self._fused_context_kv_weight.device, + ) + self._context_kv_fusion_available = True + + def _precompute_fused_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + num_ctx = context_states.shape[0] + num_layers = self._num_context_layers + + # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects + # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs). + all_kv = F.linear(context_states, self._fused_context_kv_weight) + all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width) + all_kv_c = all_kv[..., : self._context_kv_lora_rank] + all_k_pe = all_kv[..., self._context_kv_lora_rank :] + + # Layer-major layout lets the 2-D RMSNorm weights select a distinct row + # for each draft layer in one grouped kernel. + all_kv_c = all_kv_c.permute(1, 0, 2).contiguous() + all_kv_c_normed = torch.empty_like(all_kv_c) + ops.rms_norm( + all_kv_c_normed, + all_kv_c, + self._context_kv_norm_weights, + self._context_rms_norm_eps, + ) + + all_k_pe = all_k_pe.permute(1, 0, 2).contiguous() + all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim) + repeated_positions = self._context_positions_repeated[: num_layers * num_ctx] + repeated_positions.view(num_layers, num_ctx).copy_(context_positions) + # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; + # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. + rotary_emb = self.layers[0].self_attn.rotary_emb + assert rotary_emb is not None + ops.rotary_embedding( + repeated_positions, + all_k_pe_flat, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) + + if context_slot_mapping is None: + return + + cache_layers = [layer.self_attn for layer in self.layers] + if ( + not is_quantized_kv_cache(cache_layers[0].kv_cache_dtype) + and self._has_uniform_block_layout(cache_layers) + and ( + isinstance(context_slot_mapping, torch.Tensor) + or all(s is not None for s in context_slot_mapping) + ) + ): + # Grouped context KV insert only supports unquantized (bf16) KV cache + # and assumes that all layers share the same block layout. + + if isinstance(context_slot_mapping, (list, tuple)): + per_layer_slot_mappings = [ + s for s in context_slot_mapping if s is not None + ] + if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1: + # All rows alias to the same slot mapping. + slot_mapping = ( + per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1) + ) + else: + slot_mapping = torch.stack(per_layer_slot_mappings, dim=0) + else: + # Broadcast the single shared context_slot_mapping tensor. + slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1) + + ref_cache = cache_layers[0].kv_cache + ops.concat_and_cache_mla_grouped( + all_kv_c_normed, + all_k_pe.squeeze(2), + self._get_context_kv_cache_ptrs(cache_layers), + slot_mapping, + ref_cache.size(1), + ref_cache.stride(0), + ref_cache.stride(1), + ) + return + + for layer_idx, layer in enumerate(self.layers): + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn = layer.self_attn + attn.impl.do_kv_cache_update( + all_kv_c_normed[layer_idx], + all_k_pe[layer_idx], + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _has_uniform_block_layout( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> bool: + if not hasattr(self, "_layers_share_kv_block_layout"): + ref_cache = cache_layers[0].kv_cache + self._layers_share_kv_block_layout = all( + cl.kv_cache.size(1) == ref_cache.size(1) + and cl.kv_cache.stride(0) == ref_cache.stride(0) + and cl.kv_cache.stride(1) == ref_cache.stride(1) + for cl in cache_layers + ) + return self._layers_share_kv_block_layout + + def _get_context_kv_cache_ptrs( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> torch.Tensor: + # The per-layer KV cache base pointers are stable after allocation, so + # build the pointer array once and return it on every call. + if not hasattr(self, "_context_cache_ptrs"): + ref_cache = cache_layers[0].kv_cache + cache_ptrs = torch.tensor( + [cl.kv_cache.data_ptr() for cl in cache_layers], + dtype=torch.int64, + device=ref_cache.device, + ) + self._context_cache_ptrs = cache_ptrs + return self._context_cache_ptrs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + hidden_states = inputs_embeds + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + hidden_states, _ = self.final_norm(hidden_states, residual) + return hidden_states + + +class K3DSparkForCausalLM(nn.Module): + has_own_embed_tokens = False + has_own_lm_head = False + draft_id_to_target_id = None + checkpoint_skip_substrs = ("confidence_head", "embed_tokens", "lm_head") + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"": "model."}, + orig_to_new_stacked={ + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + ".q_a_proj": (".fused_qkv_a_proj", 0), + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + target_layer_num = vllm_config.model_config.get_num_layers( + vllm_config.parallel_config + ) + self.model = K3DSparkModel( + vllm_config=vllm_config, + start_layer_id=target_layer_num, + prefix=maybe_prefix(prefix, "model"), + ) + + # Assigned by load_dspark_model from the target. Keeping no placeholder + # avoids a transient full-vocabulary allocation for this 163k-vocab model. + self.lm_head: nn.Module | None = None + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.draft_vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.layer_name for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.lm_head is not None + return self.logits_processor(self.lm_head, hidden_states) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # confidence_head is training-only. The frozen target embedding and LM + # head are shared after this draft-specific checkpoint is loaded. + loader = AutoWeightsLoader( + self, + skip_substrs=list(self.checkpoint_skip_substrs), + ) + loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self.model._build_fused_context_kv_buffers() + return loaded_weights diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py new file mode 100644 index 00000000000..d4fec1e0084 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -0,0 +1,775 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch +from einops import rearrange +from torch import nn +from torch.nn.parameter import Parameter + +from vllm import _custom_ops as ops +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig +from vllm.distributed import divide, get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( + gather_initial_states, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.kimi_k3.nvidia.kda_metadata import ( + KimiK3KDAAttentionBackend, + KimiK3KDAMetadata, +) +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.v1.attention.backend import AttentionBackend + +logger = init_logger(__name__) + +_KDA_GATE_LOGBOUND_MIN = -5.0 + + +def a_log_weight_loader( + shard_axis: int, +) -> Callable[[torch.Tensor, torch.Tensor], None]: + """Load KDA A_log stored as either old 4D or current 1D weights.""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + tp_rank = get_tensor_model_parallel_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + + if loaded_weight.dim() == 4: + assert loaded_weight.shape[:2] == (1, 1), ( + f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}" + ) + assert loaded_weight.shape[-1] == 1, ( + f"Expected old A_log last dim to be 1, got {loaded_weight.shape}" + ) + loaded_weight = loaded_weight.view(loaded_weight.shape[2]) + + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + return default_weight_loader(param, loaded_weight) + + return loader + + +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks.""" + + def __init__( + self, + input_size: int, + output_sizes: list[int], + replicated_shard_id: int, + tp_size: int, + **kwargs, + ) -> None: + self.replicated_shard_id = replicated_shard_id + output_sizes = output_sizes.copy() + output_sizes[replicated_shard_id] *= tp_size + super().__init__(input_size, output_sizes, **kwargs) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + def weight_loader_v2( + self, + param: BasevLLMParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + tp_rank = self.tp_rank + param_tp_rank = getattr(param, "tp_rank", None) + if loaded_shard_id == self.replicated_shard_id: + self.tp_rank = 0 + if param_tp_rank is not None: + param.tp_rank = 0 + try: + super().weight_loader_v2(param, loaded_weight, loaded_shard_id) + finally: + self.tp_rank = tp_rank + if param_tp_rank is not None: + param.tp_rank = param_tp_rank + + +def is_fused_kda_decode_supported( + num_heads: int, + head_dim: int, + conv_width: int, + num_spec: int, + input_dtype: torch.dtype, + conv_state_dtype: torch.dtype, +) -> bool: + if ( + num_heads not in (12, 24, 48, 96) + or head_dim != 128 + or conv_width != 4 + or num_spec != 0 + or input_dtype != torch.bfloat16 + or conv_state_dtype != torch.bfloat16 + or is_conv_state_dim_first() + or not hasattr(torch.ops._C, "fused_kda_decode") + ): + return False + # SM90 is architecture-specific; SM10x and SM12x use family binaries. + return ( + current_platform.is_device_capability(90) + or current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) + + +def is_flashkda_supported( + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + capability = current_platform.get_device_capability() + return ( + capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + + +def _flashkda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + import vllm._flashkda_C # noqa: F401 + + out = torch.empty(v.shape, dtype=v.dtype, device=v.device) + final_state = torch.empty_like(initial_state) + workspace = torch.empty( + torch.ops._flashkda_C.get_workspace_size( + q.shape[0] * q.shape[1], + q.shape[2], + cu_seqlens.numel() - 1, + ), + dtype=torch.uint8, + device=q.device, + ) + # FlashKDA hardcodes dense Q/K/V/G strides. Beta may be row-strided because + # FlashKDA materializes its transposed [H, T] layout internally. + # TODO: Teach FlashKDA to consume beta in [T, H] layout directly instead + # of transposing it to contiguous [H, T] storage internally. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + q.shape[-1] ** -0.5, + out, + workspace, + A_log.contiguous(), + dt_bias.view(-1, q.shape[-1]).contiguous(), + lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + ) + return out, final_state + + +def resolve_kda_prefill_backend( + backend: str, + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, +) -> str: + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + supported = is_flashkda_supported(head_dim, dtype, lower_bound) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, " + "head_dim=128, and a bounded KDA gate." + ) + if supported and backend != "triton": + logger.info_once("Using FlashKDA KDA prefill backend.") + return "flashkda" + return "triton" + + +def _make_decode_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, + decode_conv1d_weight: torch.Tensor | None, +) -> Callable[..., None]: + sharded_dims = [dim // tp_size for dim in dims] + + def weight_loader( + param: torch.Tensor, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + if loaded_weight.dim() == 2: + loaded_weight = loaded_weight.unsqueeze(1) + shard_size = sharded_dims[loaded_shard_id] + source_start = tp_rank * shard_size + target_start = sum(sharded_dims[:loaded_shard_id]) + loaded_shard = loaded_weight[source_start : source_start + shard_size] + param.data[target_start : target_start + shard_size].copy_(loaded_shard) + if decode_conv1d_weight is not None and not param.is_meta: + decode_conv1d_weight[loaded_shard_id].copy_( + loaded_shard.squeeze(1).transpose(0, 1) + ) + + return weight_loader + + +def _make_decode_norm_weight_loader( + decode_norm_weight: torch.Tensor, +) -> Callable[..., None]: + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + default_weight_loader(param, loaded_weight) + if not param.is_meta: + decode_norm_weight.copy_(param.data) + + return weight_loader + + +class KimiK3DeltaAttention(GatedDeltaNetAttention): + def get_attn_backend(self) -> type[AttentionBackend]: + return KimiK3KDAAttentionBackend + + def get_state_dtype( + self, + ) -> tuple[torch.dtype, torch.dtype]: + if self.model_config is None or self.cache_config is None: + raise ValueError("model_config and cache_config must be set") + return MambaStateDtypeCalculator.kda_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + + def get_state_shape( + self, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + return MambaStateShapeCalculator.kda_state_shape( + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, + ) + + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__(config, vllm_config, prefix) + + kda_config = config.linear_attn_config # type: ignore[attr-defined] + assert kda_config is not None, "linear_attn_config must be set" + self.head_dim = kda_config["head_dim"] + self.num_heads = kda_config["num_heads"] + assert self.num_heads % self.tp_size == 0 + self.local_num_heads = divide(self.num_heads, self.tp_size) + self.projection_size = self.head_dim * self.num_heads + self.local_projection_size = divide(self.projection_size, self.tp_size) + self.conv_size = kda_config["short_conv_kernel_size"] + assert kda_config.get("use_full_rank_gate", False), ( + "KimiK3DeltaAttention requires a full-rank gate" + ) + + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. + qkvg_output_sizes = [self.projection_size] * 4 + in_proj_output_sizes = qkvg_output_sizes + [ + self.head_dim, + self.num_heads, + ] + local_output_size = ( + 4 * self.local_projection_size + self.head_dim + self.local_num_heads + ) + self.in_proj_padding = -local_output_size % 16 + if self.in_proj_padding: + in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( + self.hidden_size, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvgfab", + ) + if self.in_proj_padding: + self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.f_b_proj", + ) + self.dt_bias = nn.Parameter( + torch.empty(self.local_projection_size, dtype=torch.float32) + ) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + # One packed parameter and cache let decode run a single conv update. + # Prefill slices them back into Q/K/V to obtain dense outputs cheaply. + self.conv1d = ColumnParallelLinear( + input_size=self.conv_size, + output_size=3 * self.projection_size, + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.conv1d", + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + # Keep a width-major copy for fused decode without changing the layout + # consumed by the prefill and fallback decode kernels. + conv_state_dtype, _ = self.get_state_dtype() + decode_conv1d_weight = None + if is_fused_kda_decode_supported( + self.local_num_heads, + self.head_dim, + self.conv_size, + self.num_spec, + vllm_config.model_config.dtype, + conv_state_dtype, + ): + logger.info_once("Fused KDA decode kernel (conv+KDA+norm) is enabled.") + decode_conv1d_weight = torch.empty( + 3, + self.conv_size, + self.local_projection_size, + dtype=self.conv1d.weight.dtype, + device=self.conv1d.weight.device, + ) + self.register_buffer( + "decode_conv1d_weight", decode_conv1d_weight, persistent=False + ) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_decode_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + decode_conv1d_weight, + ) + }, + ) + + self.A_log = nn.Parameter( + torch.empty(self.local_num_heads, dtype=torch.float32) + ) + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) + + self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None) + if self.gate_lower_bound is not None: + assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, ( + "KDA gate lower bound must be in " + f"[{_KDA_GATE_LOGBOUND_MIN}, 0). " + f"Got {self.gate_lower_bound}." + ) + + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + self.kda_prefill_backend = resolve_kda_prefill_backend( + backend, + self.head_dim, + vllm_config.model_config.dtype, + self.gate_lower_bound, + ) + + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") + decode_norm_weight = None + if decode_conv1d_weight is not None: + decode_norm_weight = torch.empty( + self.head_dim, + dtype=torch.float32, + device=self.o_norm.weight.device, + ) + self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False) + if decode_norm_weight is not None: + # Upcast once while loading; direct BF16 norm weights slow the + # fully fused decode kernel. + if hasattr(self.o_norm.weight, "weight_loader"): + delattr(self.o_norm.weight, "weight_loader") + set_weight_attrs( + self.o_norm.weight, + {"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)}, + ) + self.o_proj = RowParallelLinear( + self.projection_size, + self.hidden_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.o_proj", + ) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden_states.size(0) + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + split_sizes = [ + 3 * self.local_projection_size, + self.local_projection_size, + self.head_dim, + self.local_num_heads, + ] + if self.in_proj_padding: + split_sizes.append(self.in_proj_padding) + projected = projected_qkvgfab.split(split_sizes, dim=-1) + mixed_qkv, g_proj_states, f_a, beta = projected[:4] + + g1 = self.f_b_proj(f_a)[0] + beta = beta.unsqueeze(0) + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) + g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, + ) + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") + return self.o_proj(core_attn_out)[0] + + @eager_break_during_capture + def _forward( + self, + mixed_qkv: torch.Tensor, + g1: torch.Tensor, + g2: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if attn_metadata_raw is None: + return + + from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + + assert isinstance(attn_metadata_raw, dict) + attn_metadata_narrowed = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata) + m = attn_metadata_narrowed + has_initial_state = m.has_initial_state + non_spec_query_start_loc = m.non_spec_query_start_loc + non_spec_state_indices_tensor = m.non_spec_state_indices_tensor + spec_token_indx = m.spec_token_indx + non_spec_token_indx = m.non_spec_token_indx + spec_state_indices_tensor = m.spec_state_indices_tensor + spec_query_start_loc = m.spec_query_start_loc + num_accepted_tokens = m.num_accepted_tokens + num_actual_tokens = m.num_actual_tokens + has_spec_decode = m.num_spec_decodes > 0 + mixed_qkv = mixed_qkv[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + beta = beta[:, :num_actual_tokens] + + conv_state, recurrent_state = self.kv_cache + # The convolution kernels consume (..., dim, width - 1). + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + + if ( + self.decode_conv1d_weight is not None + and self.decode_norm_weight is not None + and not has_spec_decode + and m.num_prefills == 0 + and m.num_decodes > 0 + ): + assert non_spec_state_indices_tensor is not None + ops.fused_kda_decode( + x=mixed_qkv, + weight=self.decode_conv1d_weight, + bias=self.conv1d.bias, + conv_state=conv_state, + raw_g=g1, + raw_beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=non_spec_state_indices_tensor[:num_actual_tokens], + state=recurrent_state, + out=core_attn_out[:, :num_actual_tokens], + lower_bound=self.gate_lower_bound, + output_gate=g2[:num_actual_tokens], + norm_weight=self.decode_norm_weight, + norm_eps=self.o_norm.eps, + ) + return + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split( + self.local_projection_size, dim=0 + ) + q_conv_state, k_conv_state, v_conv_state = conv_state.split( + self.local_projection_size, dim=-2 + ) + + # Separate multi-query speculative tokens from prefill/plain decode. + if has_spec_decode: + if m.num_prefills == 0 and m.num_decodes == 0: + mixed_qkv_spec = mixed_qkv + g1_spec, beta_spec = g1, beta + mixed_qkv_ns = g1_ns = beta_ns = None + else: + assert spec_token_indx is not None + assert non_spec_token_indx is not None + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + g1_spec = g1.index_select(1, spec_token_indx) + beta_spec = beta.index_select(1, spec_token_indx) + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + g1_ns = g1.index_select(1, non_spec_token_indx) + beta_ns = beta.index_select(1, non_spec_token_indx) + else: + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta + + # Spec-decode multi-query path. + core_attn_out_spec = None + if has_spec_decode: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] + spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_conv_out = torch.empty_like(mixed_qkv_spec) + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=spec_conv_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=spec_query_start_loc, + max_query_len=spec_max_query_len, + validate_data=False, + out=spec_conv_out, + ) + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1) + ) + spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1] + spec_out = ( + core_attn_out[:, : q_spec.shape[1]] + if m.num_prefills == 0 and m.num_decodes == 0 + else None + ) + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) + + # Prefill or plain-decode path. + core_attn_out_non_spec = None + if mixed_qkv_ns is not None: + assert g1_ns is not None and beta_ns is not None + if m.num_prefills > 0: + q_ns, k_ns, v_ns = mixed_qkv_ns.split( + self.local_projection_size, dim=-1 + ) + + # Separate convolution calls accept row-strided packed inputs + # and produce dense Q/K/V without an additional V copy. + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + q_ns, k_ns, v_ns = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) + for x in (q_ns, k_ns, v_ns) + ) + + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + initial_state = gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) + if self.kda_prefill_backend == "flashkda": + assert self.gate_lower_bound is not None + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=q_ns, + k=k_ns, + v=v_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + g_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + ) + recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + else: + # Pure non-speculative decode. + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : mixed_qkv_ns.size(0) + ] + packed_conv_out = torch.empty_like(mixed_qkv_ns) + mixed_qkv_ns = causal_conv1d_update( + mixed_qkv_ns, + conv_state, + conv_weights, + self.conv1d.bias, + activation="silu", + conv_state_indices=decode_conv_indices, + validate_data=True, + out=packed_conv_out, + ) + ( + core_attn_out_non_spec, + _, + ) = fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv_ns, + raw_g=g1_ns, + raw_beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + state_indices=decode_conv_indices, + ) + + # Restore the scheduler's original token order for mixed batches. + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + elif core_attn_out_non_spec is not None: + # TODO: prefill and decode kernels write directly to core_attn_out + core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ + 0, :num_actual_tokens + ] + else: + assert core_attn_out_spec is not None + # Triton normalizes in place, so this is a self-copy with no device + # work. Keep it for the out-of-place native implementation. + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py new file mode 100644 index 00000000000..0bb4864d291 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -0,0 +1,496 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 specialization of GDN attention metadata. + +The request classification and cudagraph staging intentionally mirror +``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its +prefill KDA kernel internally, so this builder omits the shared FLA chunk +metadata construction. +""" + +from dataclasses import dataclass +from functools import cache + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + split_decodes_and_prefills, +) +from vllm.v1.kv_cache_interface import MambaSpec + + +@cache +def _metadata_launch_pdl() -> bool: + return current_platform.is_arch_support_pdl() + + +@triton.jit(do_not_specialize=["num_requests"]) +def _get_aligned_state_indices_kernel( + block_table_ptr, + seq_lens_ptr, + state_indices_ptr, + block_table_stride_0: tl.constexpr, + block_table_stride_1: tl.constexpr, + seq_lens_stride: tl.constexpr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + num_requests, + CACHE_BLOCK_SIZE: tl.constexpr, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + valid_row = rows < num_requests + seq_lens = tl.load( + seq_lens_ptr + rows * seq_lens_stride, + mask=valid_row, + other=1, + ) + # Triton truncates signed division toward zero, unlike PyTorch floor + # division. Clamping makes both semantics equivalent for seq_lens <= 0. + first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + block_table_ptr + + rows[:, None] * block_table_stride_0 + + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + tl.store( + state_indices_ptr + + rows[:, None] * state_indices_stride_0 + + state_slots[None, :] * state_indices_stride_1, + state_indices, + mask=valid_row[:, None] & valid_state_slot[None, :], + ) + + +def _mamba_get_block_table_tensor( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + kv_cache_spec: MambaSpec, + mamba_cache_mode: str, +) -> torch.Tensor: + if mamba_cache_mode in ("all", "none"): + return block_table + + assert block_table.is_cuda and seq_lens.is_cuda + num_requests = block_table.shape[0] + num_state_slots = 1 + kv_cache_spec.num_speculative_blocks + state_indices = torch.empty( + (num_requests, num_state_slots), + dtype=block_table.dtype, + device=block_table.device, + ) + BLOCK_ROWS = 32 + grid = (triton.cdiv(num_requests, BLOCK_ROWS),) + _get_aligned_state_indices_kernel[grid]( + block_table, + seq_lens, + state_indices, + block_table.stride(0), + block_table.stride(1), + seq_lens.stride(0), + state_indices.stride(0), + state_indices.stride(1), + num_requests, + CACHE_BLOCK_SIZE=kv_cache_spec.block_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + return state_indices + + +@triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"]) +def _stage_spec_decode_metadata_kernel( + state_indices_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + staged_state_indices_ptr, + staged_query_start_loc_ptr, + staged_num_accepted_tokens_ptr, + state_indices_stride_0: tl.constexpr, + state_indices_stride_1: tl.constexpr, + staged_state_indices_stride_0: tl.constexpr, + staged_state_indices_stride_1: tl.constexpr, + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS: tl.constexpr, + BLOCK_STATE_SLOTS: tl.constexpr, + NULL_STATE_ID: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + real_request = rows < num_spec_decodes + + state_slots = tl.arange(0, BLOCK_STATE_SLOTS) + valid_state_slot = state_slots < NUM_STATE_SLOTS + state_indices = tl.load( + state_indices_ptr + + rows[:, None] * state_indices_stride_0 + + state_slots[None, :] * state_indices_stride_1, + mask=real_request[:, None] & valid_state_slot[None, :], + other=NULL_STATE_ID, + ) + tl.store( + staged_state_indices_ptr + + rows[:, None] * staged_state_indices_stride_0 + + state_slots[None, :] * staged_state_indices_stride_1, + state_indices, + mask=(rows < batch_size)[:, None] & valid_state_slot[None, :], + ) + + query_row = tl.minimum(rows, num_spec_decodes) + query_start_loc = tl.load( + query_start_loc_ptr + query_row, + mask=rows <= batch_size, + ) + tl.store( + staged_query_start_loc_ptr + rows, + query_start_loc, + mask=rows <= batch_size, + ) + + num_accepted_tokens = tl.load( + num_accepted_tokens_ptr + rows, + mask=real_request, + other=1, + ) + tl.store( + staged_num_accepted_tokens_ptr + rows, + num_accepted_tokens, + mask=rows < batch_size, + ) + + +def stage_spec_decode_metadata( + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + staged_state_indices: torch.Tensor, + staged_query_start_loc: torch.Tensor, + staged_num_accepted_tokens: torch.Tensor, + *, + num_spec_decodes: int, +) -> None: + """Stage speculative-decode metadata into CUDA-graph buffers.""" + assert state_indices.is_cuda + assert state_indices.ndim == 2 + batch_size, num_state_slots = staged_state_indices.shape + BLOCK_ROWS = 32 + grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),) + _stage_spec_decode_metadata_kernel[grid]( + state_indices, + query_start_loc, + num_accepted_tokens, + staged_state_indices, + staged_query_start_loc, + staged_num_accepted_tokens, + state_indices.stride(0), + state_indices.stride(1), + staged_state_indices.stride(0), + staged_state_indices.stride(1), + num_spec_decodes, + batch_size, + NUM_STATE_SLOTS=num_state_slots, + BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), + NULL_STATE_ID=NULL_BLOCK_ID, + BLOCK_ROWS=BLOCK_ROWS, + num_warps=1, + launch_pdl=_metadata_launch_pdl(), + ) + + +@dataclass +class KimiK3KDAMetadata(GDNAttentionMetadata): + pass + + +class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> KimiK3KDAMetadata: + m = common_attn_metadata + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + assert isinstance(self.kv_cache_spec, MambaSpec) + # Equivalent PyTorch "align" path: + # start = ((seq_lens - 1) // block_size).clamp_(min=0) + # offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32) + # indices = (start[:, None] + offsets).to(torch.int64) + # block_table_tensor = torch.gather(block_table, 1, indices) + block_table_tensor = _mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 + # A nonnegative entry identifies a spec request. If no draft token + # was scheduled, process the whole batch as non-spec instead. + if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0: + spec_sequence_masks_cpu = None + num_spec_decodes = 0 + else: + num_spec_decodes = spec_sequence_masks_cpu.sum().item() + + if num_spec_decodes == 0: + # The runner orders ordinary decodes before prefills. + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills(m, decode_threshold=1) + ) + num_spec_decode_tokens = 0 + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = None + non_spec_state_indices_tensor = block_table_tensor[:, 0] + spec_query_start_loc = None + non_spec_query_start_loc = query_start_loc + non_spec_query_start_loc_cpu = query_start_loc_cpu + num_accepted_tokens = None + else: + assert spec_sequence_masks_cpu is not None + assert num_accepted_tokens is not None + query_lens_cpu = query_start_loc_cpu.diff() + num_query_tokens = query_start_loc_cpu[-1].item() + + # Exclude zero-length cudagraph padding from request-indexed + # non-spec metadata. + active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0) + non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu] + num_non_spec_requests = non_spec_query_lens_cpu.numel() + num_non_spec_tokens = non_spec_query_lens_cpu.sum().item() + + # Query length alone cannot distinguish a true decode from a + # one-token prefill chunk. Packed decode is safe only when every + # active non-spec request is a true, one-token decode. + assert m.is_prefilling is not None + assert m.is_prefilling.device.type == "cpu" + non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu] + use_prefill = torch.any( + non_spec_is_prefilling | (non_spec_query_lens_cpu != 1) + ).item() + if use_prefill: + num_prefills = num_non_spec_requests + num_prefill_tokens = num_non_spec_tokens + num_decodes = 0 + num_decode_tokens = 0 + else: + num_prefills = 0 + num_prefill_tokens = 0 + num_decodes = num_non_spec_requests + num_decode_tokens = num_non_spec_tokens + + num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens + + if num_prefills == 0 and num_decodes == 0: + spec_token_indx = None + non_spec_token_indx = None + # Real requests precede trailing cudagraph padding. + spec_state_indices_tensor = block_table_tensor[ + :num_spec_decodes, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = None + # Padding trails real requests, so this prefix already contains + # the correct cumulative token counts. + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] + else: + query_lens = query_start_loc.diff() + spec_sequence_masks_gpu = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device + ) + spec_token_masks = torch.repeat_interleave( + spec_sequence_masks_gpu, + query_lens, + output_size=num_query_tokens, + ) + # Stable partitioning preserves request-local token order in + # both subgroup tensors. + index = torch.argsort(spec_token_masks, stable=True) + num_non_spec_tokens = num_prefill_tokens + num_decode_tokens + non_spec_token_indx = index[:num_non_spec_tokens] + spec_token_indx = index[num_non_spec_tokens:] + + # Spec requests carry one state slot per speculative step; + # non-spec requests use only their current state slot. + spec_state_indices_tensor = block_table_tensor[ + spec_sequence_masks_cpu, : self.num_spec + 1 + ] + non_spec_state_indices_tensor = block_table_tensor[ + active_non_spec_mask_cpu, 0 + ] + + spec_query_lens = query_lens[spec_sequence_masks_cpu] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + spec_query_lens, + dim=0, + out=spec_query_start_loc[1:], + ) + if num_prefills > 0: + non_spec_query_lens = query_lens[active_non_spec_mask_cpu] + non_spec_query_start_loc = torch.zeros( + non_spec_query_lens.size(0) + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + non_spec_query_lens, + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + non_spec_query_lens_cpu.size(0) + 1, + dtype=torch.int32, + ) + torch.cumsum( + non_spec_query_lens_cpu, + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + else: + # Packed decode consumes one row per request and does not + # use cumulative sequence lengths. + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + + # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares + # its own chunk indices. Only causal-convolution metadata is needed here. + nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + if num_prefills > 0: + has_initial_state = m.compute_num_computed_tokens() > 0 + if spec_sequence_masks_cpu is not None: + has_initial_state = has_initial_state[active_non_spec_mask_cpu] + assert non_spec_query_start_loc_cpu is not None + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + non_spec_query_start_loc_cpu, + device=query_start_loc.device, + ) + ) + else: + has_initial_state = None + + # Prepare per-request tensors for cudagraph replay. num_actual_tokens + # may be token-padded, while state/query/acceptance metadata is indexed + # by request. + batch_size = m.num_reqs + if ( + self.use_full_cuda_graph + and num_spec_decodes > 0 + and num_prefills == 0 + and num_decodes == 0 + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_bs + ): + # Equivalent PyTorch staging: + # state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID) + # qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1]) + # accepted[:N].copy_(accepted_src); accepted[N:].fill_(1) + stage_spec_decode_metadata( + state_indices=spec_state_indices_tensor, + query_start_loc=spec_query_start_loc, + num_accepted_tokens=num_accepted_tokens, + staged_state_indices=self.spec_state_indices_tensor[:batch_size], + staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1], + staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size], + num_spec_decodes=num_spec_decodes, + ) + + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_spec_decodes == 0 + and num_decodes <= self.decode_cudagraph_max_bs + ): + self.non_spec_state_indices_tensor[:num_decodes].copy_( + non_spec_state_indices_tensor, non_blocking=True + ) + self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_( + NULL_BLOCK_ID + ) + non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[ + :batch_size + ] + + return KimiK3KDAMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + has_initial_state=has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=None, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + ) + + +class KimiK3KDAAttentionBackend(GDNAttentionBackend): + @staticmethod + def get_name() -> str: + return "KIMI_K3_KDA" + + @staticmethod + def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]: + return KimiK3KDAMetadataBuilder diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py new file mode 100644 index 00000000000..64ea39b4875 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103. + +Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module +name plays no role. Each measured shape maps to a :class:`ProjectionSpec` +holding the winning backend per token count. The static part of the decision is +resolved once per module at install time into a small ``{M: call}`` plan, so the +per-forward path is a single dict lookup. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform + +Backend = Literal["cute", "dsv3_fused_a"] +# A resolved per-token-count call: the backend plus its CuTe config (None for +# dsv3, which needs no config). +ResolvedCall = tuple[Backend, SkinnyGemmConfig | None] + + +@dataclass(frozen=True, slots=True) +class ProjectionSpec: + n: int + k: int + dsv3_tokens: frozenset[int] = frozenset() + cute_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + residual_configs: tuple[tuple[int, SkinnyGemmConfig], ...] = () + name: str = "" # optional debug label; never used for dispatch + + def cute_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.cute_configs).get(num_tokens) + + def residual_config(self, num_tokens: int) -> SkinnyGemmConfig | None: + return dict(self.residual_configs).get(num_tokens) + + +def _cute( + num_tokens: int, + block_size: int, + outputs_per_block: int, + k_unroll: int, + vector_width: int = 8, +) -> SkinnyGemmConfig: + return SkinnyGemmConfig( + num_tokens, + block_size, + outputs_per_block, + k_unroll, + vector_width, + ) + + +_M1_TO_16 = frozenset(range(1, 17)) +_M1 = frozenset({1}) + +# Keyed by local (N, K). Where two projections share a shape (only 1536x7168: +# shared_gate_up_proj and mla_g_proj) the entry is unified. +KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = { + (1536, 128): ProjectionSpec(1536, 128, _M1_TO_16, name="f_b_proj"), + (3072, 128): ProjectionSpec(3072, 128, _M1_TO_16, name="f_b_proj"), + # 1536x7168 is shared by shared_gate_up_proj and mla_g_proj. dsv3 M1..16 is + # only crash-safe once the mla_g aux-stream/PDL capture fix lands (subtask + # task_7388aba1); the fallback if it cannot be fixed is dsv3_tokens=_M1. + (1536, 7168): ProjectionSpec( + 1536, 7168, _M1_TO_16, name="shared_gate_up_proj/mla_g_proj" + ), + (3072, 7168): ProjectionSpec( + 3072, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 3, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="shared_gate_up_proj", + ), + (2112, 7168): ProjectionSpec(2112, 7168, _M1_TO_16, name="fused_qkv_a_proj"), + (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"), + (4608, 1536): ProjectionSpec(4608, 1536, _M1_TO_16, name="q_b_proj"), + (3584, 7168): ProjectionSpec( + 3584, + 7168, + frozenset(range(2, 9)), + ((1, _cute(1, 224, 2, 4)),), + name="routed_expert_down_proj", + ), + (6288, 7168): ProjectionSpec( + 6288, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 64, 3, 2)), + (3, _cute(3, 32, 3, 4)), + (4, _cute(4, 128, 6, 1)), + ), + name="in_proj_qkvgfab", + ), + (12448, 7168): ProjectionSpec( + 12448, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + ), + name="in_proj_qkvgfab", + ), + (7168, 768): ProjectionSpec(7168, 768, _M1_TO_16, name="shared_down_proj"), + (7168, 1536): ProjectionSpec( + 7168, 1536, cute_configs=((1, _cute(1, 96, 4, 2)),), name="o_proj" + ), + (7168, 3072): ProjectionSpec( + 7168, + 3072, + cute_configs=( + (1, _cute(1, 96, 2, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="o_proj", + ), + (7168, 3584): ProjectionSpec( + 7168, + 3584, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + ), + residual_configs=( + (1, _cute(1, 64, 4, 2)), + (2, _cute(2, 64, 7, 2)), + (3, _cute(3, 64, 2, 1)), + (4, _cute(4, 64, 2, 1)), + ), + name="routed_expert_up_proj", + ), + (7168, 4224): ProjectionSpec( + 7168, + 4224, + cute_configs=((1, _cute(1, 96, 4, 2, 4)),), + name="dense_down_proj", + ), + (7168, 8448): ProjectionSpec( + 7168, + 8448, + cute_configs=( + (1, _cute(1, 32, 4, 4)), + (2, _cute(2, 96, 4, 1)), + (3, _cute(3, 96, 4, 1)), + ), + name="dense_down_proj", + ), + (8448, 7168): ProjectionSpec( + 8448, + 7168, + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (16896, 7168): ProjectionSpec( + 16896, + 7168, + cute_configs=( + (1, _cute(1, 224, 6, 4)), + (2, _cute(2, 32, 4, 4)), + ), + name="dense_gate_up_proj", + ), + (20480, 7168): ProjectionSpec( + 20480, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + (40960, 7168): ProjectionSpec( + 40960, + 7168, + cute_configs=( + (1, _cute(1, 128, 4, 2)), + (2, _cute(2, 64, 4, 2)), + (3, _cute(3, 64, 2, 2)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # TP16. Measured on B300 over M=1..16 with the same >=5% threshold as the + # entries above. The replicated projections (2112x7168, 3584x7168, + # 7168x3584) keep their shapes at TP16 and reuse the entries above, and + # o_proj lands on 7168x768, which shared_down_proj already covers. + (3216, 7168): ProjectionSpec( + 3216, + 7168, + # Both gaps in this range are measured, not oversights: dsv3 is only + # 4% ahead at M6..M8, and at M16 cuBLAS switches to a faster kernel + # (11.42us vs dsv3's 11.83us) after trailing it by 6-8% at M9..M15. + frozenset(range(9, 16)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 4, 2)), + (3, _cute(3, 128, 2, 1)), + (4, _cute(4, 64, 2, 2)), + (5, _cute(5, 128, 3, 1)), + ), + name="in_proj_qkvgfab", + ), + (768, 7168): ProjectionSpec( + 768, + 7168, + frozenset(range(5, 17)), + cute_configs=( + (1, _cute(1, 224, 2, 4)), + (2, _cute(2, 224, 2, 2)), + (3, _cute(3, 224, 2, 2)), + (4, _cute(4, 224, 2, 2)), + ), + name="mla_g_proj/shared_gate_up_proj", + ), + (1152, 1536): ProjectionSpec( + 1152, + 1536, + frozenset(range(2, 17)), + ((1, _cute(1, 192, 3, 4)),), + name="q_b_proj", + ), + (768, 128): ProjectionSpec(768, 128, _M1_TO_16, name="f_b_proj"), + # dsv3 drops under 5% from M9 on for this shape. + (7168, 384): ProjectionSpec( + 7168, 384, frozenset(range(1, 9)), name="shared_down_proj" + ), + (4224, 7168): ProjectionSpec( + 4224, + 7168, + frozenset(range(4, 9)), + cute_configs=( + (1, _cute(1, 224, 3, 4)), + (2, _cute(2, 128, 2, 1)), + (3, _cute(3, 64, 2, 2)), + ), + name="dense_gate_up_proj", + ), + (10240, 7168): ProjectionSpec( + 10240, + 7168, + cute_configs=( + (1, _cute(1, 224, 4, 2)), + (2, _cute(2, 32, 2, 4)), + (3, _cute(3, 64, 4, 1)), + (4, _cute(4, 64, 4, 1)), + ), + name="lm_head", + ), + # 7168x2112 (TP16 dense down_proj) has no entry on purpose: K=2112 divides + # none of the fused-A tile_k values, and the CuTe kernel is left with + # vector_width=2, which measured slower than cuBLAS. +} + + +def _backend_for( + spec: ProjectionSpec, num_tokens: int, has_residual: bool +) -> Backend | None: + if has_residual: + return "cute" if spec.residual_config(num_tokens) is not None else None + if spec.cute_config(num_tokens) is not None: + return "cute" + if num_tokens in spec.dsv3_tokens: + return "dsv3_fused_a" + return None + + +def select_kimi_k3_backend( + num_tokens: int, + n: int, + k: int, + *, + has_residual: bool = False, +) -> Backend | None: + """Backend for a local ``(N, K)`` at ``num_tokens``, or None to fall back.""" + spec = KIMI_K3_PROJECTIONS.get((n, k)) + return _backend_for(spec, num_tokens, has_residual) if spec is not None else None + + +def _build_plan(spec: ProjectionSpec) -> dict[int, ResolvedCall]: + plan: dict[int, ResolvedCall] = {} + for num_tokens in range(1, 17): + backend = _backend_for(spec, num_tokens, has_residual=False) + if backend == "cute": + plan[num_tokens] = ("cute", spec.cute_config(num_tokens)) + elif backend == "dsv3_fused_a": + plan[num_tokens] = ("dsv3_fused_a", None) + return plan + + +def _build_residual_plan(spec: ProjectionSpec) -> dict[int, SkinnyGemmConfig]: + return {num_tokens: config for num_tokens, config in spec.residual_configs} + + +def _is_sm103() -> bool: + return current_platform.is_device_capability((10, 3)) + + +def _is_packed_row_major(tensor: torch.Tensor) -> bool: + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) + + +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + _is_packed_row_major(x) + and _is_packed_row_major(weight) + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.is_cuda + and weight.is_cuda + and x.device == weight.device + and x.shape[1] == weight.shape[1] + ) + + +def _residual_ok(x: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor) -> bool: + return ( + residual.dim() == 2 + and residual.dtype == torch.bfloat16 + and residual.is_cuda + and residual.device == x.device + and residual.is_contiguous() + and residual.shape == (x.shape[0], weight.shape[0]) + ) + + +def _run_plan( + plan: dict[int, ResolvedCall], x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor | None: + entry = plan.get(x.shape[0]) + if entry is None: + return None + backend, config = entry + if backend == "cute": + if not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, None) + if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): + return None + output = torch.empty((x.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) + ops.dsv3_fused_a_gemm(output, x, weight.t(), enable_pdl=True) + return output + + +def _run_residual_plan( + residual_plan: dict[int, SkinnyGemmConfig], + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, +) -> torch.Tensor | None: + config = residual_plan.get(x.shape[0]) + if config is None or not shape_dynamic_skinny_gemm.is_available(): + return None + return shape_dynamic_skinny_gemm(x, weight, config, residual) + + +def try_low_latency_gemm( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor | None = None, +) -> torch.Tensor | None: + """Run the shape-selected low-latency kernel, or None to fall back. + + Resolves the plan from the shape table on each call; production installs a + precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not + use this path. + """ + if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight): + return None + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + return None + if residual is None: + return _run_plan(_build_plan(spec), x, weight) + if not _residual_ok(x, weight, residual): + return None + return _run_residual_plan(_build_residual_plan(spec), x, weight, residual) + + +class _KimiK3LowLatencyApply: + """Mixin: try the precomputed plan, else defer to the base method.""" + + def __init__(self, plan: dict[int, ResolvedCall]) -> None: + self._plan = plan + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if ( + bias is None + and not envs.VLLM_BATCH_INVARIANT + and _runtime_ok(x, layer.weight) + ): + output = _run_plan(self._plan, x, layer.weight) + if output is not None: + return output + return super().apply(layer, x, bias) # type: ignore[misc] + + +class KimiK3LowLatencyLinearMethod(_KimiK3LowLatencyApply, UnquantizedLinearMethod): + def __init__( + self, + plan: dict[int, ResolvedCall], + residual_plan: dict[int, SkinnyGemmConfig], + ) -> None: + super().__init__(plan) + self._residual_plan = residual_plan + + def apply_with_residual( + self, + layer: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + ) -> torch.Tensor: + if ( + not envs.VLLM_BATCH_INVARIANT + and _runtime_ok(x, layer.weight) + and _residual_ok(x, layer.weight, residual) + ): + output = _run_residual_plan(self._residual_plan, x, layer.weight, residual) + if output is not None: + return output + return torch.addmm(residual, x, layer.weight.t()) + + +class KimiK3LowLatencyEmbeddingMethod( + _KimiK3LowLatencyApply, UnquantizedEmbeddingMethod +): + pass + + +def enable_kimi_k3_low_latency_gemm( + module: nn.Module, + dtype: torch.dtype, +) -> None: + """Install shape-selected low-latency GEMMs and register CuTe warmups. + + Modules are matched purely by type, an exactly-unquantized method, and a + local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`. + """ + if dtype != torch.bfloat16 or not _is_sm103(): + return + + warmup_configs: set[SkinnyGemmConfig] = set() + residual_warmup_configs: set[SkinnyGemmConfig] = set() + for child in module.modules(): + is_linear = ( + isinstance(child, LinearBase) + and type(child.quant_method) is UnquantizedLinearMethod + ) + # ParallelLMHead is a VocabParallelEmbedding subclass; embed_tokens is + # the parent type, so isinstance already excludes it. + is_head = ( + isinstance(child, ParallelLMHead) + and type(child.quant_method) is UnquantizedEmbeddingMethod + ) + if not (is_linear or is_head): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2: + continue + spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1])) + if spec is None: + continue + if is_linear: + child.quant_method = KimiK3LowLatencyLinearMethod( + _build_plan(spec), _build_residual_plan(spec) + ) + else: + child.quant_method = KimiK3LowLatencyEmbeddingMethod(_build_plan(spec)) + # Warm up only the configs measured for this module's local (N, K) so a + # TP8 deployment does not compile TP4 configs and vice versa. + warmup_configs.update(config for _, config in spec.cute_configs) + residual_warmup_configs.update(config for _, config in spec.residual_configs) + + if shape_dynamic_skinny_gemm.is_available(): + if warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs) + if residual_warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs( + dtype, residual_warmup_configs, has_residual=True + ) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py new file mode 100644 index 00000000000..3334f3bbbb4 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -0,0 +1,764 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clean Multi-head Latent Attention for Kimi-K3 (NVIDIA). + +This is a self-contained MLA layer that owns the full attention path: + + hidden_states + -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) + -> explicit prefill / decode split + prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens + (+ chunked-context merge); dispatched by cache dtype + (bf16 / plain fp8 / fp8_ds_mla) + decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel + -> impl.forward_mqa -> W_UV up-proj (MQA) + -> optional output gate + -> o_proj + +Unlike ``MultiHeadLatentAttentionWrapper`` (which delegates orchestration to +``MLAAttention.forward``), this class *is* the ``AttentionLayerBase``: it selects +the backend, builds the impl, registers itself in the forward context, owns the +KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the +``DeepseekV4Attention`` structure. + +K3 specifics: optional rotary embedding (disabled for the target model's NoPE +layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``). + +Out of scope (extension points, not wired here): context parallelism (DCP/PCP), +sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. +""" + +import math +from typing import TYPE_CHECKING, cast + +import torch +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.attention import ( + _init_kv_cache_quant, + set_default_quant_scales, + should_load_quant_weights, +) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from vllm.model_executor.utils import replace_parameter +from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import ( + is_quantized_kv_cache, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionType, + MLAAttentionImpl, +) +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend +from vllm.v1.attention.ops.merge_attn_states import merge_attn_states +from vllm.v1.attention.selector import get_attn_backend +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata + +logger = init_logger(__name__) + +# Below this many tokens, overlap the g_proj GEMM on the aux stream with the +# attention front-end (the GEMM is small and launch-bound, so the overlap +# hides it); at or above it, run the gate on the main stream. +_GATE_MULTI_STREAM_TOKEN_THRESHOLD = 512 + + +@torch.compile(backend=current_platform.simple_compile_backend) +def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Apply the sigmoid output gate to a precomputed ``g_proj`` projection.""" + return attn_out * gate.sigmoid() + + +class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): + """Kimi-K3 Multi-head Latent Attention with optional RoPE and output gate.""" + + def __init__( + self, + config: KimiLinearConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + use_output_gate: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + use_rope: bool = False, + non_causal_multi_token_decode: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.non_causal_multi_token_decode = non_causal_multi_token_decode + # Latent "head" seen by the attention kernel / KV cache. + self.head_size = kv_lora_rank + qk_rope_head_dim + self.scale = self.qk_head_dim**-0.5 + self.rms_norm_eps = config.rms_norm_eps + self.layer_name = prefix + + self.rotary_emb: RotaryEmbedding | None = None + if use_rope: + rope_parameters = dict(config.rope_parameters) + if rope_parameters["rope_type"] != "default": + rope_parameters["rope_type"] = ( + "deepseek_yarn" + if rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + self.rotary_emb = get_rope( + qk_rope_head_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + dtype=torch.float32, + ) + if rope_parameters["rope_type"] == "deepseek_yarn": + mscale_all_dim = rope_parameters.get("mscale_all_dim", False) + scaling_factor = rope_parameters["factor"] + mscale = ( + 1.0 + if scaling_factor <= 1 + else 0.1 * float(mscale_all_dim) * math.log(scaling_factor) + 1.0 + ) + self.scale *= mscale * mscale + # The fused epilogues read the cos/sin table directly in fp32 and run + # the RoPE math in fp32, so there is no per-forward dtype cast (and no + # precision loss). deepseek_yarn builds cos_sin_cache in fp32 already; + # dtype=torch.float32 above forces it for the default rope too (the + # DSpark draft, which has no yarn scaling). + assert self.rotary_emb.cos_sin_cache.dtype == torch.float32, ( + "K3 fused MLA RoPE requires an fp32 cos/sin cache; got " + f"{self.rotary_emb.cos_sin_cache.dtype}." + ) + + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_heads = num_heads + self.num_local_heads = num_heads // tp_size + + # ---- Pre-attention projections (fusable front-end) ---- + # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the + # kv-down proj, or an uncompressed q path (Kimi-Linear, ``q_lora_rank`` + # None) with a standalone ``q_proj`` and separate ``kv_a_proj_with_mqa``. + if self.q_lora_rank is not None: + # Fused q-down + kv-down projection. Replicated (disable_tp) because + # the low-rank latents are shared across TP ranks; TP splitting + # happens at q_b_proj / kv_b_proj. Checkpoint weights ``q_a_proj`` + # and ``kv_a_proj_with_mqa`` map onto shards 0 and 1 respectively. + self.fused_qkv_a_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + disable_tp=True, + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + # Uncompressed query: full-rank q_proj (TP-split over heads) plus a + # replicated kv-down projection (shared latent across TP ranks). + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + + # ---- Post-attention projections ---- + self.use_output_gate = use_output_gate + self.g_proj = ( + ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_proj", + ) + if use_output_gate + else None + ) + # Aux stream (created at the model level, DeepseekV4 convention) for + # overlapping the g_proj GEMM with the attention front-end. None on + # ROCm/non-cuda -> maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self._gate_events = ( + [torch.cuda.Event(), torch.cuda.Event()] + if self.g_proj is not None and current_platform.is_cuda_alike() + else None + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # ---- Attention backend / impl / KV cache ---- + self.quant_config = quant_config + if cache_config is not None: + self.kv_cache_dtype = cache_config.cache_dtype + else: + self.kv_cache_dtype = "auto" + + dtype = torch.get_default_dtype() + self.attn_backend = get_attn_backend( + self.head_size, + dtype, + self.kv_cache_dtype, + use_mla=True, + use_sparse=False, + num_heads=self.num_local_heads, + ) + _init_kv_cache_quant(self, quant_config, prefix) + # Unit (1.0) scale for the fused fp8 prefill path: q/k/v are cast + # unscaled to match forward_mha (the prefill flash path does not + # dequantize); only the cache uses _k_scale. + self.register_buffer( + "_one_scale", torch.ones(1, dtype=torch.float32), persistent=False + ) + + impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls()) + self.impl = impl_cls( # type: ignore[assignment] + num_heads=self.num_local_heads, + head_size=self.head_size, + scale=self.scale, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=self.kv_cache_dtype, + logits_soft_cap=None, + attn_type=AttentionType.DECODER, + kv_sharing_target_layer_name=None, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + qk_head_dim=self.qk_head_dim, + v_head_dim=self.v_head_dim, + kv_b_proj=self.kv_b_proj, + indexer=None, + ) + self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) + + vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + assert ( + parallel_config.decode_context_parallel_size <= 1 + and parallel_config.prefill_context_parallel_size <= 1 + ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism." + self.prefill_backend = get_mla_prefill_backend(vllm_config)( + num_heads=self.num_local_heads, + scale=self.scale, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + vllm_config=vllm_config, + ) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # ------------------------------------------------------------------ + # AttentionLayerBase interface + # ------------------------------------------------------------------ + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + kv_cache_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + # TODO: Remove this mypy workaround once the K3 PR is fully merged. + return MLAAttentionSpec( # type: ignore[call-arg] + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=kv_cache_dtype, + cache_dtype_str=self.kv_cache_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + non_causal_multi_token_decode=self.non_causal_multi_token_decode, + ) + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + """Absorb ``kv_b_proj`` into decode-time ``W_UK_T`` / ``W_UV`` bmm weights. + + ``kv_b_proj`` produces ``[k_nope; v]`` per head from the ``kv_lora_rank`` + latent. For the MQA decode path we pre-split it so that queries are + projected into latent space by ``W_UK_T`` and the attention output is + projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. + """ + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + assert kv_b_proj_weight.shape == ( + self.kv_lora_rank, + self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim), + ), f"{kv_b_proj_weight.shape=}" + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, + self.num_local_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + W_UK, W_UV = kv_b_proj_weight.split( + [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + # (L, N, V) -> (N, L, V) + replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) + # (L, N, P) -> (N, P, L) + replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) + + quant_method = ( + self.quant_config.get_quant_method(self, prefix=self.layer_name) + if self.quant_config + else None + ) + if not should_load_quant_weights(quant_method): + set_default_quant_scales(self, register_buffer=False) + + # Precompute reciprocal scales once here (scales are final after load; + # K3 has no runtime calculate_kv_scales path) so the fp8 fused kernels + # in the decode/prefill hot path take a ready inverse instead of + # launching a per-step reciprocal kernel. + self.register_buffer( + "_q_scale_inv", self._q_scale.reciprocal().reshape(1), persistent=False + ) + self.register_buffer( + "_k_scale_inv", self._k_scale.reciprocal().reshape(1), persistent=False + ) + + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor) -> None: + """Project latent attention output back to ``v`` via ``W_UV`` (bmm).""" + # (B, N, L) -> (N, B, L) + x = x.view(-1, self.num_local_heads, self.kv_lora_rank).transpose(0, 1) + out = out.view(-1, self.num_local_heads, self.v_head_dim) + # (N, B, L) x (N, L, V) -> (N, B, V) written transposed into (B, N, V) + torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) + + def _attn_read_kv_cache(self) -> torch.Tensor: + """Latent cache as seen by the attention read kernels (decode / context). + + A plain per-tensor fp8 cache is stored as ``uint8``; view it as fp8 so + the backend reads it as E4M3 rather than fp4/E2M1 -- the latter doubles + the perceived head dim (``head_size * 2``) and fails the kernel's + ``head_dim_k == head_dim_q`` check. Mirrors ``MLAAttention.forward``; + the fp8_ds_mla layout keeps its native uint8 view. + """ + cache = self.kv_cache + if ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + ): + return cache.view(current_platform.fp8_dtype()) + return cache + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + def _forward_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Attention front-end: fused qkv-a proj -> norms -> q_b -> attention. + + Returns the pre-gate attention output ``[num_tokens, + num_local_heads * v_head_dim]``. On a profile/dummy run + it returns a zeroed buffer. + """ + if self.q_lora_rank is not None: + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] + q_c, kv_c, k_pe = qkv_lora.split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + q_c, kv_c_normed = fused_q_kv_rmsnorm( + q_c, + kv_c, + self.q_a_layernorm.weight.data, + self.kv_a_layernorm.weight.data, + self.rms_norm_eps, + ) + q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + # Uncompressed query: project directly (no q-LoRA, no q norm) and + # normalize only the kv latent. + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] + kv_c, k_pe = kv_lora.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv_c_normed = self.kv_a_layernorm(kv_c) + k_pe = k_pe.unsqueeze(1) + + attn_out = torch.empty( + (hidden_states.shape[0], self.num_local_heads * self.v_head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._attention(positions, q, kv_c_normed, k_pe, attn_out) + return attn_out + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Both branches produce (attn_out, gate); they differ only in whether + # the g_proj GEMM is overlapped on the aux stream. + g_proj = self.g_proj + events = self._gate_events + if ( + g_proj is not None + and events is not None + and self.aux_stream is not None + and hidden_states.shape[0] < _GATE_MULTI_STREAM_TOKEN_THRESHOLD + ): + attn_out, gate = maybe_execute_in_parallel( + lambda: self._forward_attn(positions, hidden_states), + lambda: g_proj(hidden_states)[0], + events[0], + events[1], + self.aux_stream, + ) + else: + attn_out = self._forward_attn(positions, hidden_states) + gate = g_proj(hidden_states)[0] if g_proj is not None else None + + if gate is not None: + attn_out = _gate_sigmoid_mul(attn_out, gate) + + # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a + # fresh private tensor, so return it directly rather than copying into a + # caller buffer -- the previous ``output[:] = ...`` convention forced an + # extra [num_tokens, hidden] copy per layer. + return self.o_proj(attn_out)[0] + + @eager_break_during_capture + def _attention( + self, + positions: torch.Tensor, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_by_layer = forward_context.attn_metadata + if attn_metadata_by_layer is None: + attn_out.zero_() + return + assert isinstance(attn_metadata_by_layer, dict) + attn_metadata = cast( + "MLACommonMetadata", attn_metadata_by_layer[self.layer_name] + ) + + num_actual_toks = attn_metadata.num_actual_tokens + slot_mapping_by_layer = forward_context.slot_mapping + assert isinstance(slot_mapping_by_layer, dict) + slot_mapping = slot_mapping_by_layer[self.layer_name] + + q = q[:num_actual_toks] + kv_c_normed = kv_c_normed[:num_actual_toks] + k_pe = k_pe[:num_actual_toks] + positions = positions[:num_actual_toks] + attn_out = attn_out[:num_actual_toks] + + cos_sin_cache = None + rope_positions = None + if self.rotary_emb is not None: + # Pass the fp32 cos/sin table straight to the fused epilogue (it reads + # fp32 and does the RoPE math in fp32) -- no per-forward dtype cast. + cos_sin_cache = self.rotary_emb.cos_sin_cache + rope_positions = positions + + # Decode tokens are laid out first, prefill tokens after. The fused + # prefill covers every supported config (bf16 / plain-fp8 / + # fp8_ds_mla), so there is no dense-MHA (forward_mha) fallback. + num_mqa_tokens = attn_metadata.num_decode_tokens + num_mha_tokens = q.size(0) - num_mqa_tokens + + # Both the prefill and decode fused epilogues write their own cache + # slice, so there is no separate do_kv_cache_update. + + # ---- Prefill: fused key-concat + cache-insert + attention ---- + if num_mha_tokens > 0: + self._forward_prefill_fused( + q[num_mqa_tokens:], + kv_c_normed[num_mqa_tokens:], + k_pe[num_mqa_tokens:], + rope_positions[num_mqa_tokens:] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[num_mqa_tokens:num_actual_toks], + attn_metadata, + attn_out[num_mqa_tokens:], + ) + + # ---- Decode: latent multi-query attention ---- + if num_mqa_tokens > 0: + mqa_q_nope, mqa_q_pe = q[:num_mqa_tokens].split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + # BMM1: absorb q_nope into latent space. (N,B,P) x (N,P,L) -> (B,N,L) + ql_nope = torch.bmm(mqa_q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1) + # Fused: concat mqa_q = [ql_nope | q_pe] and insert the decode-token + # latent into the paged cache (one launch, right before forward_mqa). + mqa_q = self._decode_concat_cache( + ql_nope, + mqa_q_pe, + kv_c_normed[:num_mqa_tokens], + k_pe[:num_mqa_tokens], + rope_positions[:num_mqa_tokens] if rope_positions is not None else None, + cos_sin_cache, + slot_mapping[:num_mqa_tokens], + ) + latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined] + mqa_q, self._attn_read_kv_cache(), attn_metadata, self + ) + self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens]) + + def _decode_concat_cache( + self, + ql_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + ) -> torch.Tensor: + """Fused decode query-concat + latent cache insert, dispatched by cache + dtype (same policy as prefill: fp8 cache -> fp8 query).""" + if self.kv_cache_dtype == "fp8_ds_mla": + cache = self.kv_cache + if cache.dtype != torch.uint8: + cache = cache.view(torch.uint8) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + ds_mla=True, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + if is_quantized_kv_cache(self.kv_cache_dtype): + assert self.impl.supports_quant_query_input, ( # type: ignore[attr-defined] + "Kimi-K3 fp8 KV cache decode requires a backend that accepts an " + "fp8 (quantized) query input." + ) + cache = self.kv_cache + if cache.dtype != torch.float8_e4m3fn: + cache = cache.view(torch.float8_e4m3fn) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + cache, + slot_mapping, + q_scale_inv=self._q_scale_inv, + cache_scale_inv=self._k_scale_inv, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + return fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + self.kv_cache, + slot_mapping, + positions=positions, + cos_sin_cache=cos_sin_cache, + ) + + def _forward_prefill_fused( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + slot_mapping: torch.Tensor, + attn_metadata, + out: torch.Tensor, + ) -> None: + """Prefill using the fused key-concat + cache-insert kernel. + + Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one + fused kernel launch, dispatched by cache dtype. The chunked context + gather + online-softmax merge are delegated to the impl. + + Supported configs (K3 fp8 policy): + - bf16 cache -> bf16 prefill query + - plain fp8 cache -> fp8 prefill query (unscaled q/k/v; cache _k_scale) + - fp8_ds_mla cache -> bf16 prefill query (656B per-tile self-scaled) + """ + prefill = attn_metadata.prefill + has_context = prefill.chunked_context is not None + fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype() + + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + if self.kv_cache_dtype == "fp8_ds_mla": + # fp8_ds_mla cache (656B, per-tile self-scaled); bf16 attention. + assert not fp8_prefill, ( + "Kimi-K3 fp8_ds_mla uses a bf16 prefill query; fp8 prefill " + "query is not supported with fp8_ds_mla." + ) + kv_cache = self.kv_cache + if kv_cache.dtype != torch.uint8: + kv_cache = kv_cache.view(torch.uint8) + k = fused_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + elif is_quantized_kv_cache(self.kv_cache_dtype): + assert fp8_prefill, ( + "Kimi-K3 fp8 KV cache requires an fp8 prefill query; enable " + "--attention-config '{\"use_prefill_query_quantization\": true}'." + ) + # Plain per-tensor fp8: quant q/k/v (unscaled, matching forward_mha's + # unscaled `.to(fp8)`) and insert the fp8 latent (scaled by _k_scale). + kv_cache = self.kv_cache + if kv_cache.dtype != torch.float8_e4m3fn: + kv_cache = kv_cache.view(torch.float8_e4m3fn) + q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + kv_cache, + slot_mapping, + self._one_scale, + self._one_scale, + self._one_scale, + self._k_scale_inv, + positions, + cos_sin_cache, + ) + else: + # Concat full K = [k_nope | k_pe] and insert [kv_c_normed | k_pe] + # into the paged cache for these prefill tokens, in one launch. + k = fused_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + self.kv_cache, + slot_mapping, + positions, + cos_sin_cache, + ) + + # When there is no chunked context, backends that honor `out` write the + # attention result straight into it, avoiding a slice+flatten+copy. + writes_out = not has_context and prefill.prefill_backend.supports_out() + output_prefill = prefill.prefill_backend.run_prefill_new_tokens( + q=q, + k=k, + v=v, + return_softmax_lse=has_context, + out=( + out.view(-1, self.num_local_heads, self.v_head_dim) + if writes_out + else None + ), + ) + + if has_context: + context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] + q, self._attn_read_kv_cache(), attn_metadata, self._k_scale + ) + suffix_output, suffix_lse = output_prefill + out = out.view(-1, self.num_local_heads, self.v_head_dim) + merge_attn_states( + output=out, + prefix_output=context_output[..., : self.v_head_dim], + prefix_lse=context_lse, + suffix_output=suffix_output[..., : self.v_head_dim], + suffix_lse=suffix_lse, + prefill_tokens_with_context=prefill.chunked_context.prefill_tokens_with_context, + ) + elif not writes_out: + out.copy_(output_prefill[..., : self.v_head_dim].flatten(start_dim=-2)) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py new file mode 100644 index 00000000000..ec2dc651fac --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -0,0 +1,1859 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal model implementation for vLLM.""" + +import math +from collections.abc import Iterable +from typing import Any, cast + +import torch +from torch import nn + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear +from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + fused_grouped_topk, +) +from vllm.model_executor.layers.fused_moe.runner.latent_moe_runner import ( + LatentMoERunner, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import ( + KimiGatedDeltaNetAttention as KimiLinearGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors import ( + compressed_tensors, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + SupportsEagle3, + SupportsEncoderCudaGraph, + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) +from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs +from vllm.model_executor.models.kimi_k25_vit import ( + KimiK25MultiModalProjector, + MoonViT3dPretrainedModel, + vision_tower_forward, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + init_vllm_registered_model, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts +from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention +from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( + enable_kimi_k3_low_latency_gemm, +) +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.ops import attn_res +from vllm.models.kimi_k3.nvidia.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_reduce_scatter, + sp_shard, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import NestedTensors +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.math_utils import cdiv +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream +from vllm.v1.worker.ubatching import dbo_current_ubatch_id + +from ..common.mm_preprocess import ( + KimiK3DummyInputsBuilder, + KimiK3MultiModalProcessor, + KimiK3ProcessingInfo, +) + +logger = init_logger(__name__) + + +class KimiMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + use_sequence_parallel: bool = False, + prefix: str = "", + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=use_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=use_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act == "silu": + self.act_fn = SiluAndMul() + elif hidden_act == "situ": + self.act_fn = SituAndMul( + beta=activation_situ_beta or 1.0, + linear_beta=activation_situ_linear_beta, + ) + else: + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu and situ are supported." + ) + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class KimiRoutedOutputTransform(nn.Module): + def __init__( + self, + norm: RMSNorm | None, + up_proj: ReplicatedLinear, + ) -> None: + super().__init__() + self.norm = norm + self.up_proj = up_proj + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.norm is not None: + hidden_states = self.norm(hidden_states) + hidden_states, _ = self.up_proj(hidden_states) + return hidden_states + + +class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): + """Kimi K3 adapter for the DeepGEMM MegaMoE kernel.""" + + _kimi_symm_buffer_cache: dict[tuple[object, ...], object] = {} + _synchronized_ep_groups: set[tuple[int, int]] = set() + + def __init__( + self, + *args, + activation: str, + activation_beta: float | None, + activation_linear_beta: float | None, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.activation = activation + self.activation_beta = activation_beta + self.activation_linear_beta = activation_linear_beta + + def synchronize_first_launch(self) -> None: + ep_group = get_ep_group() + device = torch.accelerator.current_device_index() + key = (id(ep_group.cpu_group), device) + if key in self._synchronized_ep_groups: + return + torch.accelerator.synchronize() + torch.distributed.barrier(group=ep_group.cpu_group) + self._synchronized_ep_groups.add(key) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, 32), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, 32), + self.num_local_experts, + ) + self._transformed_l1_weights, self._transformed_l2_weights = ( + deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + activation=self.activation, + ) + ) + self.w13_weight = None + self.w13_weight_scale = None + self.w2_weight = None + self.w2_weight_scale = None + + def get_symm_buffer(self): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + self.activation, + ) + symm_buffer = self._kimi_symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + activation=self.activation, + ) + self._kimi_symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + self.synchronize_first_launch() + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"Kimi K3 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but its symmetric buffer supports {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + is_padding = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + if is_padding is not None: + is_padding = is_padding[:num_tokens] + + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + if is_padding is not None: + topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, + ) + + prepare_megamoe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + is_padding=is_padding, + ) + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + activation=self.activation, + activation_beta=self.activation_beta, + activation_linear_beta=self.activation_linear_beta, + fast_math=fast_math, + ) + return y + + +def make_kimi_k3_mega_moe_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + mapping = [] + for expert_id in range(num_experts): + for shard_id in ("w1", "w2", "w3"): + param_prefix = "w13" if shard_id in ("w1", "w3") else "w2" + for suffix in ("weight_packed", "weight_scale"): + param_suffix = "weight" if suffix == "weight_packed" else suffix + mapping.append( + ( + f"experts.{param_prefix}_{param_suffix}", + f"experts.{expert_id}.{shard_id}.{suffix}", + expert_id, + shard_id, + ) + ) + return mapping + + +class KimiMoE(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + use_sequence_parallel: bool = False, + ): + super().__init__() + hidden_size = config.hidden_size + moe_intermediate_size = config.moe_intermediate_size + num_experts = config.num_experts + num_experts_per_token = config.num_experts_per_token + assert moe_intermediate_size is not None + assert num_experts is not None + assert num_experts_per_token is not None + moe_renormalize = config.moe_renormalize + routed_expert_hidden_size = config.routed_expert_hidden_size + self.use_latent_moe = routed_expert_hidden_size is not None + self.moe_hidden_size = ( + routed_expert_hidden_size + if routed_expert_hidden_size is not None + else hidden_size + ) + self.latent_moe_use_norm = config.latent_moe_use_norm + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_renormalize = moe_renormalize + self.use_grouped_topk = config.use_grouped_topk + self.num_expert_group = config.num_expert_group + self.topk_group = config.topk_group + self.moe_router_activation_func = config.moe_router_activation_func + self.num_shared_experts = config.num_shared_experts + self.layer_idx = layer_idx + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "Kimi K3 MegaMoE requires expert parallel. Enable it with " + "--enable-expert-parallel." + ) + if self.use_mega_moe and config.hidden_act != "situ": + raise ValueError("Kimi K3 MegaMoE requires SITU activation.") + if self.use_mega_moe and not self.use_latent_moe: + raise ValueError("Kimi K3 MegaMoE requires latent MoE projections.") + if self.use_mega_moe and not self.use_grouped_topk: + raise ValueError("Kimi K3 MegaMoE requires grouped top-k routing.") + if self.use_mega_moe and (self.num_expert_group != 1 or self.topk_group != 1): + raise NotImplementedError( + "Kimi K3 MegaMoE currently requires one expert group." + ) + self.padded_moe_intermediate_size = moe_intermediate_size + min_moe_intermediate_per_partition = getattr( + config, "min_moe_intermediate_per_partition", 256 + ) + if self.tp_size > 1: + moe_intermediate_per_partition = moe_intermediate_size // self.tp_size + if moe_intermediate_per_partition < min_moe_intermediate_per_partition: + self.padded_moe_intermediate_size = ( + min_moe_intermediate_per_partition * self.tp_size + ) + activation_situ_beta = ( + config.activation_situ_beta if config.hidden_act == "situ" else None + ) + activation_situ_linear_beta = ( + config.activation_situ_linear_beta if config.hidden_act == "situ" else None + ) + + # Route with fp32 logits for numerically stable expert selection. + self.gate = GateLinear( + input_size=hidden_size, + output_size=num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=torch.float32) + ) + + if self.num_shared_experts is not None: + shared_intermediate_size = moe_intermediate_size * self.num_shared_experts + self.shared_experts = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + use_sequence_parallel=use_sequence_parallel, + prefix=f"{prefix}.shared_experts", + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + ) + else: + self.shared_experts = None + + self.routed_expert_down_proj: ReplicatedLinear | None + self.routed_expert_norm: RMSNorm | None + self.routed_expert_up_proj: ReplicatedLinear | None + self.routed_output_transform: KimiRoutedOutputTransform | None + if self.use_latent_moe: + self.routed_expert_down_proj = ReplicatedLinear( + hidden_size, + self.moe_hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_down_proj", + ) + self.routed_expert_norm = ( + RMSNorm(self.moe_hidden_size, eps=config.rms_norm_eps) + if self.latent_moe_use_norm + else None + ) + # Replicated up-proj: the full weight lives on every rank and + # produces the full hidden dim locally. This lets LatentMoERunner + # fuse the latent and shared reductions into a single all-reduce + # (concat the two partials, reduce once), then run the up-proj and + # shared add locally with no further collective. + self.routed_expert_up_proj = ReplicatedLinear( + self.moe_hidden_size, + hidden_size, + bias=False, + quant_config=None, + prefix=f"{prefix}.routed_expert_up_proj", + ) + + self.routed_output_transform = KimiRoutedOutputTransform( + self.routed_expert_norm, self.routed_expert_up_proj + ) + # Auxiliary CUDA stream to overlap the router gate with the routed + # down projection on decode-sized batches (gated by + # VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD). + self._down_proj_stream: torch.cuda.Stream | None = aux_stream() + self._down_proj_events = (torch.cuda.Event(), torch.cuda.Event()) + else: + self.routed_expert_down_proj = None + self.routed_expert_norm = None + self.routed_expert_up_proj = None + self.routed_output_transform = None + + if self.use_mega_moe: + ep_group = get_ep_group() + ep_size = ep_group.world_size + ep_rank = ep_group.rank_in_group + if num_experts % ep_size != 0: + raise ValueError( + f"Kimi K3 num_experts={num_experts} must be divisible by " + f"EP size {ep_size}." + ) + num_local_experts = num_experts // ep_size + self.experts = KimiK3MegaMoEExperts( + vllm_config, + num_experts=num_experts, + num_local_experts=num_local_experts, + experts_start_idx=ep_rank * num_local_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + prefix=f"{prefix}.experts", + activation="situ", + activation_beta=activation_situ_beta, + activation_linear_beta=activation_situ_linear_beta, + ) + else: + enable_tail_fusion = envs.VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION + self.experts = FusedMoE( + shared_experts=self.shared_experts, + num_experts=num_experts, + top_k=num_experts_per_token, + hidden_size=self.moe_hidden_size, + intermediate_size=self.padded_moe_intermediate_size, + activation=config.hidden_act, + activation_situ_beta=activation_situ_beta, + activation_situ_linear_beta=activation_situ_linear_beta, + renormalize=moe_renormalize, + quant_config=quant_config, + use_grouped_topk=config.use_grouped_topk, + num_expert_group=config.num_expert_group, + topk_group=config.topk_group, + prefix=f"{prefix}.experts", + scoring_func=config.moe_router_activation_func, + e_score_correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + # Down projection runs outside FusedMoE so it can overlap the + # router gate on the aux stream (see forward()); the original + # hidden states are passed to forward() as shared_experts_input + # so shared experts still see the untransformed input. + routed_input_transform=None, + routed_output_transform=self.routed_output_transform, + is_sequence_parallel=use_sequence_parallel, + runner_cls=LatentMoERunner if self.use_latent_moe else None, + runner_args=( + {"enable_k3_latent_moe_tail_fusion": enable_tail_fusion} + if self.use_latent_moe + else None + ), + ) + if self.padded_moe_intermediate_size != moe_intermediate_size: + w13_weight = getattr(self.experts, "w13_weight", None) + if w13_weight is None: + w13_weight = getattr(self.experts, "w13_weight_packed", None) + w2_weight = getattr(self.experts, "w2_weight", None) + if w2_weight is None: + w2_weight = getattr(self.experts, "w2_weight_packed", None) + if w13_weight is not None: + w13_weight.data.zero_() + if w2_weight is not None: + w2_weight.data.zero_() + self.experts.moe_config.intermediate_size_per_partition_unpadded = ( + moe_intermediate_size // self.tp_size + ) + + def _maybe_overlap_router_and_down_proj( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Compute the routed-expert down projection alongside the router, + overlapping them on separate CUDA streams when latent MoE is enabled. + + The router gate and the down projection both read ``hidden_states``, so + the gate runs on the default stream and the down projection on the aux + stream, joined via ``maybe_execute_in_parallel``. For MegaMoE the + grouped top-k selection consumes only the gate logits, so it also runs + on the default stream and overlaps the down projection. + + Returns: + ``(routed_hidden_states, router_output, topk_ids)``. + ``routed_hidden_states`` is the down-projected latent (or the + original ``hidden_states`` when latent MoE is disabled). For MegaMoE + ``router_output`` holds the grouped top-k weights and ``topk_ids`` + the selected experts; otherwise ``router_output`` holds the raw gate + logits and ``topk_ids`` is ``None``. + """ + + def _router( + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + router_logits, _ = self.gate(hidden_states) + if not self.use_mega_moe: + return router_logits, None + return fused_grouped_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.experts.top_k, + renormalize=self.moe_renormalize, + e_score_correction_bias=self.gate.e_score_correction_bias.data, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + scoring_func=self.moe_router_activation_func, + routed_scaling_factor=self.routed_scaling_factor, + ) + + down_proj = self.routed_expert_down_proj + if down_proj is None: + router_output, topk_ids = _router(hidden_states) + return hidden_states, router_output, topk_ids + num_tokens = hidden_states.shape[0] + (router_output, topk_ids), (routed_hidden_states, _) = ( + maybe_execute_in_parallel( + lambda: _router(hidden_states), + lambda: down_proj(hidden_states), + self._down_proj_events[0], + self._down_proj_events[1], + self._down_proj_stream + if num_tokens <= envs.VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD + else None, + ) + ) + return routed_hidden_states, router_output, topk_ids + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + # Overlap the gate with the routed down projection; the returned hidden + # states are already down-projected. Keep the original ``hidden_states`` + # for the shared experts. + routed_hidden_states, router_output, topk_ids = ( + self._maybe_overlap_router_and_down_proj(hidden_states) + ) + if self.use_mega_moe: + assert self.routed_output_transform is not None + assert topk_ids is not None + final_hidden_states = self.experts( + routed_hidden_states, + router_output, + topk_ids, + activation_clamp=None, + ) + final_hidden_states = self.routed_output_transform(final_hidden_states) + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts( + hidden_states + ) + else: + # Routed experts consume the down-projected latent; shared experts + # (inside FusedMoE) get the original hidden states via + # shared_experts_input. + final_hidden_states = self.experts( + hidden_states=routed_hidden_states, + router_logits=router_output, + shared_experts_input=hidden_states, + ) + return final_hidden_states.view(num_tokens, hidden_size) + + +class KimiDecoderLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = int(prefix.rsplit(".", 1)[1]) + + self.is_moe = config.is_moe + layer_idx = self.layer_idx + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + self.is_moe_layer = ( + self.is_moe + and config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + if config.is_kda_layer(layer_idx): + kda_config = config.linear_attn_config + assert kda_config is not None + # This class also serves standalone Kimi-Linear through the model + # registry. Only Kimi-K3's full-rank gate uses the private KDA path. + if kda_config.get("use_full_rank_gate", False): + self.self_attn = KimiK3DeltaAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = False + else: + self.self_attn = KimiLinearGatedDeltaNetAttention( + config, + vllm_config, + prefix=f"{prefix}.self_attn", + ) + self._self_attn_writes_output = True + else: + qk_nope_head_dim = config.qk_nope_head_dim + qk_rope_head_dim = config.qk_rope_head_dim + v_head_dim = config.v_head_dim + kv_lora_rank = config.kv_lora_rank + mla_use_nope = config.mla_use_nope + assert qk_nope_head_dim is not None + assert qk_rope_head_dim is not None + assert v_head_dim is not None + assert kv_lora_rank is not None + assert mla_use_nope, "Kimi-K3 MLA (MultiHeadLatentAttention) is NoPE-only" + # q_lora_rank may be None (Kimi-Linear): the MLA layer then uses an + # uncompressed q_proj instead of the fused q-LoRA front-end. + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=kv_lora_rank, + use_output_gate=bool(config.mla_use_output_gate), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + aux_stream=aux_stream, + ) + self._self_attn_writes_output = False + + if self.use_sequence_parallel: + self.self_attn.o_proj.reduce_results = False + + if self.is_moe_layer: + self.block_sparse_moe = KimiMoE( + config=config, + vllm_config=vllm_config, + quant_config=quant_config, + prefix=f"{prefix}.block_sparse_moe", + layer_idx=layer_idx, + use_sequence_parallel=self.use_sequence_parallel, + ) + self.mlp = self.block_sparse_moe + else: + self.mlp = KimiMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + use_sequence_parallel=self.use_sequence_parallel, + activation_situ_beta=config.activation_situ_beta, + activation_situ_linear_beta=config.activation_situ_linear_beta, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + attn_res_block_size = config.attn_res_block_size + self.use_attn_res = attn_res_block_size is not None + if self.use_attn_res: + assert attn_res_block_size is not None + self.attn_res_block_size = attn_res_block_size + self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 + self.block_write_idx = layer_idx // self.attn_res_block_size + self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) + self.self_attention_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.self_attention_res_proj", + ) + self.mlp_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.mlp_res_proj", + ) + + def _run_self_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if self._self_attn_writes_output: + output = torch.empty_like(hidden_states) + self.self_attn( + hidden_states=hidden_states, + positions=positions, + output=output, + ) + return output + return self.self_attn( + hidden_states=hidden_states, + positions=positions, + ) + + def _pre_attn_norm( + self, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + assert hidden_states is not None + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + assert residual is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.self_attention_res_norm.weight, + self.self_attention_res_proj.weight.squeeze(0), + self.input_layernorm.weight, + num_blocks=self.prev_valid_blocks, + 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, + ) + return hidden_states, prefix_sum, residual + + def _post_attn_norm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + prefix_sum: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if not self.use_attn_res: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + return hidden_states, prefix_sum, residual + + assert prefix_sum is not None + if self.is_block_write_layer: + prefix_sum = hidden_states + prefix_delta = None + else: + prefix_delta = hidden_states + mlp_valid_blocks = self.prev_valid_blocks + self.is_block_write_layer + hidden_states = attn_res( + prefix_sum, + prefix_delta, + residual, + self.mlp_res_norm.weight, + self.mlp_res_proj.weight.squeeze(0), + self.post_attention_layernorm.weight, + num_blocks=mlp_valid_blocks, + block_write_idx=-1, + eps=self.mlp_res_norm.variance_epsilon, + output_norm_eps=self.post_attention_layernorm.variance_epsilon, + ) + return hidden_states, prefix_sum, residual + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor | None, + residual: torch.Tensor | None, + prefix_sum: 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 + ) + assert hidden_states is not None + + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states) + # Remove SP padding before attention. + hidden_states = hidden_states[: positions.shape[0]] + + # Attention. + hidden_states = self._run_self_attn(positions, hidden_states) + + if self.use_sequence_parallel: + # Add SP padding if needed, and then perform reduce scatter. + hidden_states = sp_reduce_scatter(hidden_states) + + hidden_states, prefix_sum, residual = self._post_attn_norm( + hidden_states, residual, prefix_sum + ) + + # MoE/MLP. + hidden_states = self.mlp(hidden_states) + return hidden_states, prefix_sum, residual + + +class KimiLinearModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + 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 + parallel_config = vllm_config.parallel_config + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + self.use_sequence_parallel = ( + parallel_config.pipeline_parallel_size == 1 + and parallel_config.enable_expert_parallel + and parallel_config.tensor_parallel_size > 1 + and (use_mega_moe or parallel_config.data_parallel_size > 1) + ) + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + # Aux stream for overlapping the MLA g_proj output-gate GEMM with the + # attention front-end (DeepseekV4 convention: created at the model + # level and threaded into each attention layer). + aux_stream = torch.cuda.Stream() + + def get_layer(prefix: str): + return KimiDecoderLayer( + config, + vllm_config, + prefix, + aux_stream=aux_stream, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + self.num_attn_res_blocks = ( + cdiv(self.end_layer, self.attn_res_block_size) + if self.attn_res_block_size is not None + else 0 + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.use_attn_res: + self.output_attn_res_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=f"{prefix}.output_attn_res_proj", + ) + else: + self.norm = PPMissingLayer() + if self.use_attn_res: + self.output_attn_res_norm = PPMissingLayer() + self.output_attn_res_proj = PPMissingLayer() + + world_size = get_tensor_model_parallel_world_size() + assert config.num_attention_heads % world_size == 0, ( + "num_attention_heads must be divisible by world_size" + ) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + residual_shape: tuple[int, ...] = (batch_size, self.config.hidden_size) + if self.use_attn_res: + assert self.attn_res_block_size is not None + residual_shape = ( + batch_size, + cdiv(self.start_layer, self.attn_res_block_size), + self.config.hidden_size, + ) + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ), + "residual": torch.zeros(residual_shape, dtype=dtype, device=device), + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + assert hidden_states is not None + + aux_hidden_states: list[torch.Tensor] = [] + if self.start_layer in self.aux_hidden_state_layers: + if self.use_attn_res or residual is None: + aux_hidden_states.append(hidden_states) + else: + aux_hidden_states.append(hidden_states + residual) + + full_num_tokens = positions.shape[0] + if self.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + assert residual is None, "Currently, SP is not supported with PP" + + 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), + ) + if residual is not None: + block_residual[:, : residual.size(1), :].copy_(residual) + prefix_sum = hidden_states + hidden_states = None + residual = block_residual + + for layer_idx, layer in enumerate( + self.layers[self.start_layer : self.end_layer], + start=self.start_layer, + ): + hidden_states, prefix_sum, residual = layer( + positions=positions, + hidden_states=hidden_states, + prefix_sum=prefix_sum, + residual=residual, + ) + if (layer_idx + 1) in self.aux_hidden_state_layers: + if self.use_attn_res: + assert prefix_sum is not None + aux_hidden_state = prefix_sum + hidden_states + else: + assert residual is not None + aux_hidden_state = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded aux hidden states. + # TODO: Optimize this. + aux_hidden_state = sp_all_gather(aux_hidden_state) + aux_hidden_state = aux_hidden_state[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + + assert hidden_states is not None + assert residual is not None + if not get_pp_group().is_last_rank: + assert not self.use_sequence_parallel, ( + "Currently, SP is not supported with PP" + ) + if prefix_sum is not None: + hidden_states = hidden_states + prefix_sum + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if self.use_attn_res: + assert prefix_sum is not None + hidden_states = attn_res( + prefix_sum, + hidden_states, + residual, + self.output_attn_res_norm.weight, + self.output_attn_res_proj.weight.squeeze(0), + None, + num_blocks=self.num_attn_res_blocks, + block_write_idx=-1, + eps=self.output_attn_res_norm.variance_epsilon, + output_norm_eps=0.0, + ) + else: + hidden_states = hidden_states + residual + + if self.use_sequence_parallel: + # Gather SP-sharded hidden states. + hidden_states = sp_all_gather(hidden_states) + hidden_states = hidden_states[:full_num_tokens] + + # NOTE: the final norm is applied in compute_logits instead of here, so + # the MTP draft model receives the pre-norm hidden states. + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights( + self, + weights: Iterable[ + tuple[str, torch.Tensor] | tuple[str, torch.Tensor, dict[str, Any]] + ], + ) -> set[str]: + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + use_mega_moe = any( + module.use_mega_moe + for module in self.modules() + if isinstance(module, KimiMoE) + ) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping( + self.config.num_experts + ) + elif self.config.is_moe: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any( + n.endswith("w13_weight_packed") for n in params_dict + ) + loaded_params: set[str] = set() + for args in weights: + name, loaded_weight = args[0], args[1] + kwargs: dict[str, Any] = args[2] if len(args) > 2 else {} + if "rotary_emb.inv_freq" in name: + continue + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is not None: + continue # skip spec decode layers for main model + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name_mapped = name.replace(weight_name, param_name) + # Packed projections are only present on compatible layers. + if name_mapped not in params_dict: + continue + name = name_mapped + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name = name.replace(expert_weight_name, expert_param_name) + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + expert_id=expert_id, + shard_id=expert_shard_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") + and name not in params_dict + and not self.config.is_linear_attn + ): # noqa: E501 + continue + # Remapping the name of FP8 kv-scale. + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, **kwargs) + loaded_params.add(name) + return loaded_params + + def finalize_mega_moe_weights(self) -> None: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + +class KimiLinearForCausalLM( + nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid +): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model_config = vllm_config.model_config + self.vllm_config = vllm_config + self.config = self.model_config.hf_config + quant_config = vllm_config.quant_config + self.quant_config = quant_config + self.model = KimiLinearModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + enable_kimi_k3_low_latency_gemm(self, self.model_config.dtype) + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + return self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.kda_state_dtype( + vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.kda_state_shape( + tp_size, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], + num_spec=num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.kda_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # The model's final norm is applied here (not at the end of forward) so + # that the pre-norm hidden states can be fed to the MTP draft model. + hidden_states = self.model.norm(hidden_states, None) + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + ) + loaded = loader.load_weights(weights) + self.model.finalize_mega_moe_weights() + # The fused MultiHeadLatentAttention's process_weights_after_loading + # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load + # hook for any AttentionLayerBase, so no manual trigger is needed here. + return loaded + + +def get_spec_layer_idx_from_weight_name( + config: KimiLinearConfig, weight_name: str +) -> int | None: + if hasattr(config, "num_nextn_predict_layers") and ( + config.num_nextn_predict_layers > 0 + ): + layer_idx = config.num_hidden_layers + for i in range(config.num_nextn_predict_layers): + # Match regardless of the surrounding prefix. The name may arrive as + # ``model.layers.{i}.``, a bare ``layers.{i}.`` (after AutoWeightsLoader + # has stripped the ``model.`` prefix in the main model), or with the + # multimodal ``language_model.model.layers.{i}.`` prefix. + if f"layers.{layer_idx + i}." in weight_name: + return layer_idx + i + return None + + +@MULTIMODAL_REGISTRY.register_processor( + KimiK3MultiModalProcessor, + info=KimiK3ProcessingInfo, + dummy_inputs=KimiK3DummyInputsBuilder, +) +class KimiK3ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsQuant, + SupportsEagle3, + HasInnerState, + IsHybrid, +): + """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" + + supports_encoder_tp_data = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "mm_projector.proj.0": "mm_projector.linear_1", + "mm_projector.proj.2": "mm_projector.linear_2", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "<|kimi_image_placeholder|>" + raise ValueError(f"Unsupported modality: {modality}") + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + config: KimiK3Config = model_config.hf_config + self.config = config + self.model_config = model_config + quant_config = vllm_config.quant_config + + multimodal_config = model_config.multimodal_config + assert multimodal_config is not None + self.use_data_parallel = is_vit_use_data_parallel( + config.vision_config.num_attention_heads + ) + self.hidden_size = config.text_config.hidden_size + self.device = current_platform.current_device() + + with self._mark_tower_model(vllm_config, "image"): + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "vision_tower"), + ) + if self._maybe_ignore_quant_config(quant_config) is not None: + self.vision_tower = self.vision_tower.to(device=self.device) + else: + self.vision_tower = self.vision_tower.to( + device=self.device, dtype=model_config.dtype + ) + + vision_attn = self.vision_tower.encoder.blocks[0].attn + if vision_attn.is_flash_attn_backend and vision_attn._fa_version == 4: + from vllm.models.kimi_k3.nvidia.ops.vision_fa4_warmup import ( + KimiK3VisionFA4WarmupConfig, + register_kimi_k3_vision_fa4_warmup, + ) + + merge_height, merge_width = config.vision_config.merge_kernel_size + mm_config = model_config.get_multimodal_config() + assert mm_config is not None + register_kimi_k3_vision_fa4_warmup( + KimiK3VisionFA4WarmupConfig( + num_heads=vision_attn.num_heads, + head_dim=vision_attn.head_size, + dtype=vision_attn.dtype, + max_batch_size=( + vllm_config.scheduler_config.max_num_seqs + * mm_config.get_limit_per_prompt("image") + ), + max_seqlen=( + vllm_config.scheduler_config.max_num_encoder_input_tokens + * merge_height + * merge_width + ), + ) + ) + + self.mm_projector = KimiK25MultiModalProjector( + config=config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=self._maybe_ignore_quant_config(quant_config), + prefix=maybe_prefix(prefix, "mm_projector"), + ) + self.mm_projector = self.mm_projector.to( + device=self.device, dtype=model_config.dtype + ) + + self.quant_config = quant_config + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["KimiLinearForCausalLM"], + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + self.media_placeholder: int = self.config.media_placeholder_token_id + + # -- SupportsEncoderCudaGraph protocol methods -- + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "pixel_values", + "pos_embeds", + "rope_freqs_cis", + "cu_seqlens", + "max_seqlen", + "sequence_lengths", + "merge_gather_idx", + ], + out_hidden_size=self.hidden_size, + ) + + def get_encoder_cudagraph_budget_range( + self, vllm_config: VllmConfig + ) -> tuple[int, int]: + min_budget = 64 + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return min_budget, max_budget + + @staticmethod + def _get_grid_thws(mm_kwargs: dict[str, Any]) -> list[list[int]]: + grid_thws = mm_kwargs["grid_thws"] + if not isinstance(grid_thws, list): + grid_thws = grid_thws.tolist() + return grid_thws + + @staticmethod + def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor: + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + pixel_values = torch.cat(pixel_values) + if pixel_values.ndim in (3, 5): + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], + *pixel_values.shape[2:], + ) + return pixel_values + + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + kh, kw = self.config.vision_config.merge_kernel_size + return [ + EncoderItemSpec( + input_size=t * h * w, + output_tokens=(h // kh) * (w // kw), + ) + for t, h, w in self._get_grid_thws(mm_kwargs) + ] + + def select_encoder_cudagraph_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + grid_thws = self._get_grid_thws(mm_kwargs) + pixel_values = self._get_pixel_values(mm_kwargs) + source_grid = mm_kwargs["grid_thws"] + + if not indices: + empty_grid = ( + source_grid[:0] if isinstance(source_grid, torch.Tensor) else [] + ) + return {"pixel_values": pixel_values[:0], "grid_thws": empty_grid} + + patch_counts = [t * h * w for t, h, w in grid_thws] + offsets = [0] + for count in patch_counts: + offsets.append(offsets[-1] + count) + selected_pixel_values = torch.cat( + [pixel_values[offsets[i] : offsets[i + 1]] for i in indices] + ) + grid_device = ( + source_grid.device if isinstance(source_grid, torch.Tensor) else None + ) + selected_grid = torch.tensor( + [grid_thws[i] for i in indices], + dtype=torch.long, + device=grid_device, + ) + return {"pixel_values": selected_pixel_values, "grid_thws": selected_grid} + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + kh, kw = self.config.vision_config.merge_kernel_size + per_item_output = (token_budget + max_batch_size - 1) // max_batch_size + rope = self.vision_tower.encoder.rope_2d + max_output_width = rope.max_width // kw + max_output_height = rope.max_height // kh + output_width = min(math.ceil(math.sqrt(per_item_output)), max_output_width) + output_height = (per_item_output + output_width - 1) // output_width + if output_height > max_output_height: + output_height = max_output_height + output_width = (per_item_output + output_height - 1) // output_height + if output_width > max_output_width: + raise ValueError( + f"Encoder CUDA graph budget {token_budget} exceeds K3 RoPE " + f"capacity for max_batch_size={max_batch_size}" + ) + grid_thws = [ + [1, output_height * kh, output_width * kw] for _ in range(max_batch_size) + ] + + patch_size: int | tuple[int, int] = self.config.vision_config.patch_size + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + total_patches = sum(t * h * w for t, h, w in grid_thws) + pixel_values = torch.randn( + total_patches, + 3, + patch_size[0], + patch_size[1], + device=device, + dtype=dtype, + ) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + grid_thws, + max_batch_size=max_batch_size, + max_seqlen_override=max( + token_budget * kh * kw, + max(t * h * w for t, h, w in grid_thws), + ), + device=device, + ) + return EncoderCudaGraphCaptureInputs( + values=metadata | {"pixel_values": pixel_values} + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + pixel_values = self._get_pixel_values(mm_kwargs) + metadata = self.vision_tower.prepare_encoder_cudagraph_metadata( + self._get_grid_thws(mm_kwargs), + max_batch_size=max_batch_size, + device=pixel_values.device, + ) + return EncoderCudaGraphReplayBuffers( + values=metadata | {"pixel_values": pixel_values} + ) + + def _project_encoder_features(self, image_features: torch.Tensor) -> torch.Tensor: + projector_dtype = next(self.mm_projector.parameters()).dtype + if image_features.dtype != projector_dtype: + image_features = image_features.to(projector_dtype) + output = self.mm_projector(image_features) + return output.reshape(-1, output.shape[-1]) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + pixel_values = values.pop("pixel_values") + image_features = self.vision_tower(pixel_values, None, encoder_metadata=values) + return self._project_encoder_features(image_features) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + image_features = self.vision_tower( + self._get_pixel_values(mm_kwargs).to( + next(self.vision_tower.parameters()).dtype + ), + self._get_grid_thws(mm_kwargs), + ) + return self._project_encoder_features(torch.cat(image_features)) + + def _maybe_ignore_quant_config( + self, quant_config: QuantizationConfig | None + ) -> QuantizationConfig | None: + if isinstance(quant_config, compressed_tensors.CompressedTensorsConfig): + return None + return quant_config + + def _parse_and_validate_media_input( + self, **kwargs: object + ) -> KimiK25MediaPixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + grid_thws = kwargs.pop("grid_thws", None) + if pixel_values is None: + return None + + if isinstance(pixel_values, list): + pixel_values = torch.cat(cast(list[torch.Tensor], pixel_values), dim=0) + if not isinstance(pixel_values, torch.Tensor): + raise TypeError( + "pixel_values must be a tensor or a list of tensors, " + f"got {type(pixel_values)}" + ) + + if len(pixel_values.shape) == 5 or len(pixel_values.shape) == 3: + pixel_values = pixel_values.reshape( + pixel_values.shape[0] * pixel_values.shape[1], *pixel_values.shape[2:] + ) + + target_dtype = next(self.vision_tower.parameters()).dtype + pixel_values = pixel_values.to(target_dtype) + assert isinstance(grid_thws, torch.Tensor), ( + f"expect grid_thws to be a tensor, got {type(grid_thws)}" + ) + grid_thws = grid_thws.reshape(-1, grid_thws.shape[-1]) + assert grid_thws.ndim == 2 and grid_thws.size(1) == 3, ( + f"unexpected shape for grid_thws: {grid_thws.shape}" + ) + + return KimiK25MediaPixelInputs( + type="pixel_values", + pixel_values=pixel_values, + grid_thws=grid_thws, + ) + + def _process_media_input( + self, media_input: KimiK25MediaPixelInputs + ) -> list[torch.Tensor]: + media_features = vision_tower_forward( + self.vision_tower, + media_input["pixel_values"], + media_input["grid_thws"], + mm_projector=self.mm_projector, + use_data_parallel=self.use_data_parallel, + ) + return media_features + + def embed_multimodal(self, **kwargs: object) -> NestedTensors | None: + media_input = self._parse_and_validate_media_input(**kwargs) + if media_input is None: + return None + return self._process_media_input(media_input) + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if intermediate_tensors is not None: + inputs_embeds = None + return self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): + return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( + input_buffers, **kwargs + ) + + def get_seqlen_agnostic_capture_inputs(self, batch_size: int): + return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( + batch_size + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) + return KimiLinearForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) + + @classmethod + def get_mamba_state_copy_func(cls): + return KimiLinearForCausalLM.get_mamba_state_copy_func() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/models/kimi_k3/nvidia/mtp.py b/vllm/models/kimi_k3/nvidia/mtp.py new file mode 100644 index 00000000000..f08d9304b44 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/mtp.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Kimi-K3 Multi-Token-Prediction (MTP) draft model.""" + +import copy +from collections.abc import Iterable + +import torch +import torch.nn as nn + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +from ..common.mtp import fused_mtp_input +from .low_latency_gemm import enable_kimi_k3_low_latency_gemm +from .model import ( + KimiDecoderLayer, + KimiMoE, + get_spec_layer_idx_from_weight_name, + make_kimi_k3_mega_moe_expert_params_mapping, +) +from .ops.sequence_parallel import sp_all_gather, sp_padding_mask, sp_shard + +logger = init_logger(__name__) + + +class SharedHead(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class KimiK3MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: KimiLinearConfig, + vllm_config: VllmConfig, + prefix: str, + ) -> None: + super().__init__() + self.config = config + quant_config = vllm_config.quant_config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=quant_config + ) + # The MTP block starts without the base model's AttnRes state. + block_config = copy.copy(config) + block_config.attn_res_block_size = None + # NOTE: the prefix must end in the numeric spec-layer index so that + # KimiDecoderLayer can parse ``layer_idx`` and pick MLA (full attn). + # Own aux stream for the MLA g_proj output-gate overlap (DeepseekV4 + # convention: the MTP block creates its own stream). + aux_stream = torch.cuda.Stream() + self.mtp_block = KimiDecoderLayer( + block_config, vllm_config, prefix=prefix, aux_stream=aux_stream + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert inputs_embeds is not None + hidden_states = self.eh_proj( + fused_mtp_input( + positions, + inputs_embeds, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + ) + if self.mtp_block.use_sequence_parallel: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + forward_context.is_padding = sp_padding_mask( + forward_context.is_padding, hidden_states + ) + hidden_states = sp_shard(hidden_states) + + hidden_states, _, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + residual=None, + ) + if self.mtp_block.use_sequence_parallel: + assert residual is not None + hidden_states = hidden_states + residual + hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] + logits_hidden_states = self.shared_head.norm(hidden_states) + return logits_hidden_states, hidden_states + + # Produce the normalized logits input and the pre-norm recurrent state + # in one fused add-RMSNorm launch. + logits_hidden_states, hidden_states = self.shared_head.norm( + hidden_states, residual + ) + return logits_hidden_states, hidden_states + + +class KimiK3MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: KimiLinearConfig = vllm_config.model_config.hf_text_config + self.config = config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + + self.layers = torch.nn.ModuleDict( + { + str(idx): KimiK3MultiTokenPredictorLayer( + config, vllm_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) + return logits + + +class KimiK3MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_text_config + self.quant_config = vllm_config.quant_config + self.model = KimiK3MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + enable_kimi_k3_low_latency_gemm(self, vllm_config.model_config.dtype) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Mirror KimiLinearForCausalLM.load_weights naming: leading-dot shard + # names, q_lora-conditional fused QKV, and w1/w2/w3 expert weights. + kda_config = self.config.linear_attn_config + use_full_rank_gate = bool( + kda_config and kda_config.get("use_full_rank_gate", False) + ) + beta_shard_id = 5 if use_full_rank_gate else 3 + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".in_proj_qkvgfab", ".q_proj", 0), + (".in_proj_qkvgfab", ".k_proj", 1), + (".in_proj_qkvgfab", ".v_proj", 2), + (".in_proj_qkvgfab", ".b_proj", beta_shard_id), + (".in_proj_qkvgfab", ".f_a_proj", 4), + (".conv1d", ".q_conv1d", 0), + (".conv1d", ".k_conv1d", 1), + (".conv1d", ".v_conv1d", 2), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + if use_full_rank_gate: + stacked_params_mapping.append((".in_proj_qkvgfab", ".g_proj", 3)) + if getattr(self.config, "q_lora_rank", None) is not None: + stacked_params_mapping += [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + ] + + use_mega_moe = any( + module.use_mega_moe + for module in self.modules() + if isinstance(module, KimiMoE) + ) + if self.config.is_moe and use_mega_moe: + expert_params_mapping = make_kimi_k3_mega_moe_expert_params_mapping( + self.config.num_experts + ) + elif self.config.is_moe: + expert_params_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_experts, + ) + else: + expert_params_mapping = [] + + pp_missing_layer_names = get_pp_missing_layer_names(self) + params_dict = dict(self.named_parameters()) + # Under the MXFP4 quant interface the routed experts register unpacked + # params (``w13_weight``), while the compressed-tensors checkpoint names + # them ``.weight_packed``. Rebind so the expert mapping resolves; scales + # already share the ``.weight_scale`` suffix. + experts_unpacked = not use_mega_moe and not any( + n.endswith("w13_weight_packed") for n in params_dict + ) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + # The multimodal checkpoint prefixes text weights with + # ``language_model.``; strip it so names match this draft model's + # parameter paths (``model.layers.{i}.``). Non-text weights + # (vision_tower, mm_projector, ...) never match a spec layer below. + if name.startswith("language_model."): + name = name[len("language_model.") :] + if experts_unpacked and name.endswith(".weight_packed"): + name = name.replace(".weight_packed", ".weight") + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + continue + name = self._rewrite_spec_layer_name(spec_layer, name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + # Routed experts (``.experts.{i}.w1/w2/w3``) are handled by the + # expert mapping below; skip them here. Shared experts + # (``.shared_experts.``) use gate/up_proj and fall through. + if ".experts." in name: + continue + name_mapped = name.replace(weight_name, param_name) + # Only take this mapping if the fused destination actually + # exists (e.g. QKV fusion is only present when q_lora is used). + if name_mapped not in params_dict: + continue + if name_mapped in pp_missing_layer_names: + continue + name = name_mapped + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for ( + expert_param_name, + expert_weight_name, + expert_id, + expert_shard_id, + ) in expert_params_mapping: + if expert_weight_name not in name: + continue + name_mapped = name.replace(expert_weight_name, expert_param_name) + if name_mapped in pp_missing_layer_names: + continue + param = params_dict[name_mapped] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is None: + continue + name = remapped_name + + # The embedding is shared across MTP layers; only the first + # spec layer carries the hoisted (non-".layers") copy. + if spec_layer != self.model.mtp_start_layer_idx and ( + ".layers" not in name + ): + continue + if name in pp_missing_layer_names: + continue + # The base model uses an attn-residual scheme whose per-layer + # weights (self_attention_res_*, mlp_res_*) are not used by + # the draft block; such names have no matching parameter and + # are safely skipped. + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + # Validate that weights were loaded for each expected MTP layer. + loaded_layers: set[int] = set() + for param_name in loaded_params: + spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name) + if spec_layer is not None: + loaded_layers.add(spec_layer) + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_layers: + raise ValueError( + f"MTP speculative decoding layer {layer_idx} weights " + f"missing from checkpoint. The checkpoint may not include " + f"the MTP layer weights. Use a checkpoint that includes " + f"MTP layer weights, or disable speculative decoding." + ) + + if use_mega_moe: + for module in self.modules(): + if isinstance(module, KimiMoE) and module.use_mega_moe: + module.experts.finalize_weights() + + return loaded_params + + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: + """Rewrite a checkpoint weight name to this module's parameter path. + + Top-level MTP submodules (enorm/hnorm/eh_proj/shared_head) stay under + ``model.layers.{spec_layer}.*``; the shared ``embed_tokens`` is hoisted + to ``model.*``; everything else is a transformer-block weight and gets + ``.mtp_block`` inserted. + """ + spec_layer_weight_names = [ + "embed_tokens", + "enorm", + "hnorm", + "eh_proj", + "shared_head", + ] + shared_weight_names = ["embed_tokens"] + spec_layer_weight = False + shared_weight = False + for weight_name in spec_layer_weight_names: + if weight_name in name: + spec_layer_weight = True + if weight_name in shared_weight_names: + shared_weight = True + break + if not spec_layer_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py new file mode 100644 index 00000000000..7eda574b83e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CuTe DSL kernels for KimiK3LatentMoETailOp.""" + +from .allreduce_rmsnorm_reduce_scatter_early_exit import CollectiveKernel +from .fused_add_multicast_gemm import AdaptiveUpProjectionKernel +from .lamport_copy import LamportCopyKernel + +__all__ = [ + "AdaptiveUpProjectionKernel", + "CollectiveKernel", + "LamportCopyKernel", +] diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py new file mode 100644 index 00000000000..4f33f892064 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -0,0 +1,1012 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Specialized from FlashInfer's oneshotAllreduceFusionKernel. + +"""Routed AllReduce/RMSNorm with CTA-specialized ReduceScatter early exit.""" + +from __future__ import annotations + +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from .primitives import ( + NUM_LAMPORT_BUFFERS, + PACKED_BYTES, + VEC_BF16, + bf16x8_to_packed_u32x4, + block_sum_specialized, + fragment_is_dirty, + load_global_u32x4, + load_volatile_u32, + map_shared_to_peer, + packed_u32x4_to_bf16x8, + red_async_release_gpu_add_u32, + sanitize_negative_zero, + store_global_u32x4, + store_lamport_sentinel_128, + store_shared_cluster_f32, + to_cute, + to_cute_dynamic_m, +) + + +def _mapping( + tp_size: int, + latent_dim: int, + hidden_dim: int, +) -> tuple[int, int, int, int]: + shard_dim = hidden_dim // tp_size + cluster_ctas = latent_dim // shard_dim + threads = shard_dim // VEC_BF16 + shared_roles = tp_size // cluster_ctas + return shard_dim, cluster_ctas, threads, shared_roles + + +def validate_shape(*, tp_size: int, latent_dim: int, hidden_dim: int) -> None: + """Validate constraints imposed by the fused collective mapping.""" + + if tp_size <= 0 or latent_dim <= 0 or hidden_dim <= 0: + raise ValueError("collective dimensions must be positive") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by tp_size") + shard, cluster, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + if shard % VEC_BF16: + raise ValueError("hidden_dim / tp_size must be divisible by 8") + if latent_dim % shard: + raise ValueError("latent_dim must be an integer multiple of shard_dim") + if cluster > 16 or cluster & (cluster - 1): + raise ValueError("collective cluster width must be a power of two <= 16") + if tp_size % cluster: + raise ValueError("tp_size must be divisible by collective cluster width") + if not 32 <= threads <= 1024: + raise ValueError("collective threads per CTA must be in [32, 1024]") + if threads < cluster: + raise ValueError("collective threads must cover every cluster CTA") + + +def _select_routed_schedule( + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, +) -> tuple[int, int]: + """Match upstream MNNVL's one-cluster-per-token occupancy policy.""" + + _, cluster_ctas, threads, _ = _mapping(tp_size, latent_dim, hidden_dim) + sm_count = torch.cuda.get_device_properties( + torch.accelerator.current_device_index() + ).multi_processor_count + while max_m * cluster_ctas > sm_count and cluster_ctas > 1 and threads <= 512: + cluster_ctas //= 2 + threads *= 2 + return threads, cluster_ctas + + +class AllReduceRMSNormWithReduceScatterEarlyExit: + """One routed role plus one ReduceScatter role per destination group.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool = False, + include_reduce_scatter: bool = True, + include_routed: bool = True, + ): + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + if not 0 <= rank < tp_size: + raise ValueError(f"rank must be in [0,{tp_size}), got {rank}") + if not include_routed and not include_reduce_scatter: + raise ValueError("at least one collective role must be enabled") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + ( + self.shard_dim, + mapped_cluster, + mapped_threads, + shared_roles, + ) = _mapping(tp_size, latent_dim, hidden_dim) + if include_reduce_scatter: + self.cluster_ctas = mapped_cluster + self.threads = mapped_threads + else: + self.threads, self.cluster_ctas = _select_routed_schedule( + tp_size, latent_dim, hidden_dim, max_m + ) + # Diagnostic specialization: a one-role grid executes only the routed + # AllReduce/RMSNorm path. Keep this compile-time so the production + # fused path is unchanged when include_reduce_scatter=True. + if include_routed and include_reduce_scatter: + self.roles = 1 + shared_roles + elif include_routed: + self.roles = 1 + else: + self.roles = shared_roles + self.warps = (self.threads + 31) // 32 + self.last_warp_lanes = self.threads - (self.warps - 1) * 32 + self.last_warp_mask = (1 << self.last_warp_lanes) - 1 + # The upstream MNNVL oneshot protocol assigns one cluster to each + # token. Keep that exact ownership in the routed-only diagnostic; + # reusing a cluster for multiple token waves allows the Lamport + # generation metadata to change between waves. + self.token_ctas = ( + min(max_m, max_token_ctas) if include_reduce_scatter else max_m + ) + self.fp32_internal = fp32_internal + self.include_reduce_scatter = include_reduce_scatter + self.include_routed = include_routed + + @cute.jit + def __call__( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + stream: cuda.CUstream, + ): + self.kernel( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + ).launch( + grid=(self.token_ctas, self.cluster_ctas, self.roles), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_ctas, 1), + smem=(self.warps + self.cluster_ctas) * 4, + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + token_cta, cta_y, role = cute.arch.block_idx() + logical_role = role + if cutlass.const_expr(not self.include_routed): + # A shared-only grid starts at z=0, while _token_device reserves + # logical role 0 for the routed collective. + logical_role = role + Int32(1) + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + + cute.arch.griddepcontrol_wait() + token = token_cta + while token < m: + self._token_device( + latent_source, + gamma, + latent_output, + routed_workspace, + latent_flags, + latent_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + m, + epsilon, + token, + token_cta, + cta_y, + logical_role, + cluster_rank, + tidx, + ) + token = token + self.token_ctas + + @cute.jit + def _token_device( + self, + latent_source: cute.Tensor, + gamma: cute.Tensor, + latent_output: cute.Tensor, + routed_workspace: cute.Tensor, + latent_flags: cute.Tensor, + latent_multicast_ptr: Int64, + shared_source: cute.Tensor, + shared_output: cute.Tensor, + shared_workspace: cute.Tensor, + shared_flags: cute.Tensor, + shared_peer_ptrs: cute.Tensor, + m: Int32, + epsilon: Float32, + token: Int32, + token_cta: Int32, + cta_y: Int32, + role: Int32, + cluster_rank: Int32, + tidx: Int32, + ): + if role == 0: + # ---------------- routed AllReduce + RMSNorm ---------------- + packed_idx = cluster_rank * self.threads + tidx + element_offset = ( + Int64(token) * self.latent_dim + Int64(packed_idx) * VEC_BF16 + ) + current_index = cute.arch.load((latent_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((latent_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load( + (latent_flags.iterator + 2).llvm_ptr, Uint32 + ) + dirty_num_stages = cute.arch.load( + (latent_flags.iterator + 3).llvm_ptr, Uint32 + ) + bytes_to_clear = cute.arch.load( + (latent_flags.iterator + 4).llvm_ptr, Uint32 + ) + current_elements = Int64(current_index) * ( + Int64(bytes_per_buffer) // Int64(2) + ) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + local_ptr = cute.make_ptr( + BFloat16, + (latent_source.iterator + element_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero( + load_global_u32x4(local_ptr, volatile=False) + ) + multicast_offset = ( + Int64(current_index) * Int64(bytes_per_buffer) + + ( + (Int64(token) * self.tp_size + self.rank) * self.latent_dim + + Int64(packed_idx) * VEC_BF16 + ) + * 2 + ) + store_global_u32x4( + latent_multicast_ptr + multicast_offset, + local_packed, + volatile=False, + ) + + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(latent_flags.iterator + 8, Uint32(1)) + + global_tid = ( + Int64(token) * self.cluster_ctas + Int64(cta_y) + ) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.cluster_ctas * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + ( + routed_workspace.iterator + + dirty_elements + + clear_idx * VEC_BF16 + ).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + rank_words = cute.make_rmem_tensor( + cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32 + ) + for word in cutlass.range_constexpr(4): + rank_words[self.rank, word] = local_packed[word] + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + if cutlass.const_expr(source_rank != self.rank): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) + * self.latent_dim + + Int64(packed_idx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (routed_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8( + rank_words[source_rank, None].load() + ).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + + # Preserve the original early PDL point before RMSNorm. + cute.arch.griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.fp32_internal): + # High-precision fused mode: retain the rank reduction in + # FP32 through the RMS square and row reduction. + norm_input = accum.load() + norm_square = norm_input * norm_input + else: + norm_input_bf16 = accum.load().to(BFloat16) + norm_input = norm_input_bf16.to(Float32) + # Upstream-compatible mode: FlashInfer evaluates BF16 * + # BF16 first, then promotes the rounded square to FP32. + norm_square = (norm_input_bf16 * norm_input_bf16).to(Float32) + thread_sum = norm_square.reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_tensor( + Float32, cute.make_layout((self.warps,)), byte_alignment=4 + ) + cluster_sums = smem.allocate_tensor( + Float32, + cute.make_layout((self.cluster_ctas,)), + byte_alignment=4, + ) + block_sum = block_sum_specialized( + thread_sum, + warp_sums, + tidx, + self.warps, + self.last_warp_lanes, + self.last_warp_mask, + ) + if tidx < self.cluster_ctas: + local_slot = cluster_sums.iterator + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, block_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_ctas): + full_sum = full_sum + cluster_sums[peer] + inv_rms = cute.math.rsqrt( + full_sum / Float32(self.latent_dim) + epsilon, fastmath=True + ) + gamma_ptr = cute.make_ptr( + BFloat16, + (gamma.iterator + Int64(packed_idx) * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_values = packed_u32x4_to_bf16x8( + load_global_u32x4(gamma_ptr, volatile=False) + ) + result = (norm_input * inv_rms * gamma_values.to(Float32)).to(BFloat16) + store_global_u32x4( + Int64((latent_output.iterator + element_offset).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + # The x=0 CTA rotates only after reaching its final token wave. + # Waiting for all M arrivals then guarantees every token-wave CTA + # loaded the current generation before the metadata is advanced. + if ( + token_cta == 0 + and token + self.token_ctas >= m + and cta_y == 0 + and tidx == 0 + ): + access_counter = latent_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + while arrived < Uint32(m): + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2) + cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (latent_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((latent_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((latent_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (latent_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + else: + # ---------------- shared ReduceScatter ---------------- + shared_group = role - 1 + destination = shared_group * self.cluster_ctas + cluster_rank + current_index = cute.arch.load((shared_flags.iterator + 0).llvm_ptr, Uint32) + dirty_index = cute.arch.load((shared_flags.iterator + 1).llvm_ptr, Uint32) + bytes_per_buffer = cute.arch.load( + (shared_flags.iterator + 2).llvm_ptr, Uint32 + ) + dirty_num_stages = cute.arch.load( + (shared_flags.iterator + 3).llvm_ptr, Uint32 + ) + bytes_to_clear = cute.arch.load( + (shared_flags.iterator + 4).llvm_ptr, Uint32 + ) + current_elements = Int64(current_index) * ( + Int64(bytes_per_buffer) // Int64(2) + ) + dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2)) + + source_element = ( + Int64(token) * self.hidden_dim + + Int64(destination) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + source_ptr = cute.make_ptr( + BFloat16, + (shared_source.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + local_packed = sanitize_negative_zero( + load_global_u32x4(source_ptr, volatile=False) + ) + peer_base = cute.arch.load( + (shared_peer_ptrs.iterator + destination).llvm_ptr, + Int64, + ) + destination_element = current_elements + ( + (Int64(token) * self.tp_size + self.rank) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + store_global_u32x4( + peer_base + destination_element * 2, + local_packed, + volatile=False, + ) + + # One arrival per shared destination group and token. + cute.arch.cluster_arrive() + if cluster_rank == 0 and tidx < 32: + cute.arch.cluster_wait() + if tidx == 0: + red_async_release_gpu_add_u32(shared_flags.iterator + 8, Uint32(1)) + + global_tid = ( + Int64(token) * self.tp_size + Int64(destination) + ) * self.threads + Int64(tidx) + total_threads = Int64(m) * self.tp_size * self.threads + clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES + clear_idx = global_tid + if dirty_num_stages > Uint32(0): + while clear_idx < clear_fragments: + clear_ptr = cute.make_ptr( + BFloat16, + ( + shared_workspace.iterator + + dirty_elements + + clear_idx * VEC_BF16 + ).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(clear_ptr) + clear_idx = clear_idx + total_threads + + if destination == self.rank: + rank_words = cute.make_rmem_tensor( + cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32 + ) + valid = False + while not valid: + valid = True + for source_rank in cutlass.range_constexpr(self.tp_size): + remote_element = current_elements + ( + (Int64(token) * self.tp_size + source_rank) * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + remote_ptr = cute.make_ptr( + BFloat16, + (shared_workspace.iterator + remote_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + remote = load_global_u32x4(remote_ptr, volatile=True) + for word in cutlass.range_constexpr(4): + rank_words[source_rank, word] = remote[word] + valid = valid & (not fragment_is_dirty(remote)) + + accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = Float32(0.0) + for source_rank in cutlass.range_constexpr(self.tp_size): + values = packed_u32x4_to_bf16x8( + rank_words[source_rank, None].load() + ).to(Float32) + for element in cutlass.range_constexpr(VEC_BF16): + accum[element] = accum[element] + values[element] + result = accum.load().to(BFloat16) + output_element = ( + Int64(token) * self.hidden_dim + + self.rank * self.shard_dim + + Int64(tidx) * VEC_BF16 + ) + store_global_u32x4( + Int64((shared_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(result), + volatile=False, + ) + + if destination == self.rank: + cute.arch.barrier() + + cute.arch.griddepcontrol_launch_dependents() + + if ( + token_cta == 0 + and token + self.token_ctas >= m + and shared_group == 0 + and cluster_rank == 0 + and tidx == 0 + ): + access_counter = shared_flags.iterator + 8 + arrived = load_volatile_u32(access_counter) + target = Uint32(m) * Uint32(self.tp_size // self.cluster_ctas) + while arrived < target: + arrived = load_volatile_u32(access_counter) + next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2) + cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index) + cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index) + cute.arch.store( + (shared_flags.iterator + 2).llvm_ptr, + bytes_per_buffer, + ) + cute.arch.store((shared_flags.iterator + 3).llvm_ptr, Uint32(1)) + cute.arch.store((shared_flags.iterator + 4).llvm_ptr, actual_bytes) + for index in cutlass.range_constexpr(5, 8): + cute.arch.store( + (shared_flags.iterator + index).llvm_ptr, + Uint32(0), + ) + cute.arch.store(access_counter.llvm_ptr, Uint32(0)) + + +_COMPILED: dict[tuple[object, ...], Any] = {} + + +def _routed_workspace_cute(workspace: torch.Tensor): + return to_cute(workspace.view(torch.bfloat16), 16) + + +def _compile_key( + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool, + include_routed: bool, +): + return ( + torch.accelerator.current_device_index(), + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + + +def _runtime_args( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, +): + return ( + to_cute_dynamic_m(latent_source, mode=0, assumed_align=16), + to_cute(gamma, 16), + to_cute(latent_output, 16), + _routed_workspace_cute(routed_workspace), + to_cute(routed_flags, 16), + Int64(routed_multicast_ptr), + to_cute_dynamic_m(shared_source, mode=0, assumed_align=16), + to_cute(shared_output, 16), + to_cute(shared_workspace, 16), + to_cute(shared_flags, 16), + to_cute(shared_peer_ptrs, 16), + Int32(latent_source.shape[0]), + Float32(rms_eps), + cuda.CUstream(torch.cuda.current_stream(latent_source.device).cuda_stream), + ) + + +def compile_kernel( + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + """Compile the rank/M specialization without retaining caller tensors.""" + + key = _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + if key in _COMPILED: + return + device = latent_output.device + latent = torch.empty((max_m, latent_dim), dtype=torch.bfloat16, device=device) + gamma = torch.empty((latent_dim,), dtype=torch.bfloat16, device=device) + shared = torch.empty((max_m, hidden_dim), dtype=torch.bfloat16, device=device) + kernel = AllReduceRMSNormWithReduceScatterEarlyExit( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[key] = cute.compile( + kernel, + *_runtime_args( + latent, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ), + ) + + +def launch( + latent_source: torch.Tensor, + gamma: torch.Tensor, + latent_output: torch.Tensor, + routed_workspace: torch.Tensor, + routed_flags: torch.Tensor, + routed_multicast_ptr: int, + shared_source: torch.Tensor, + shared_output: torch.Tensor, + shared_workspace: torch.Tensor, + shared_flags: torch.Tensor, + shared_peer_ptrs: torch.Tensor, + rms_eps: float, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + fp32_internal: bool, + include_reduce_scatter: bool = True, + include_routed: bool = True, +) -> None: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=latent_output, + routed_workspace=routed_workspace, + routed_flags=routed_flags, + routed_multicast_ptr=routed_multicast_ptr, + shared_output=shared_output, + shared_workspace=shared_workspace, + shared_flags=shared_flags, + shared_peer_ptrs=shared_peer_ptrs, + rms_eps=rms_eps, + fp32_internal=fp32_internal, + include_reduce_scatter=include_reduce_scatter, + include_routed=include_routed, + ) + _COMPILED[ + _compile_key( + rank, + tp_size, + latent_dim, + hidden_dim, + max_m, + max_token_ctas, + fp32_internal, + include_reduce_scatter, + include_routed, + ) + ]( + *_runtime_args( + latent_source, + gamma, + latent_output, + routed_workspace, + routed_flags, + routed_multicast_ptr, + shared_source, + shared_output, + shared_workspace, + shared_flags, + shared_peer_ptrs, + rms_eps, + ) + ) + + +class CollectiveKernel: + """Own and launch the routed AllReduce/RMSNorm plus shared ReduceScatter.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + max_token_ctas: int, + rms_eps: float, + fp32_internal: bool, + ) -> None: + validate_shape( + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + ) + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.max_token_ctas = max_token_ctas + self.rms_eps = float(rms_eps) + self.fp32_internal = fp32_internal + device = torch.device("cuda", torch.accelerator.current_device_index()) + + bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2 + routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer + self._routed_workspace = symm_mem.empty( + routed_bytes // 4, + dtype=torch.float32, + device=device, + ) + self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group) + self._routed_workspace.fill_(-0.0) + actual_bytes_per_buffer = ( + self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16 + ) + if actual_bytes_per_buffer < bytes_per_routed_buffer: + raise RuntimeError("routed symmetric workspace is too small") + self._routed_flags = torch.tensor( + [0, 2, actual_bytes_per_buffer, 0, 0, 0, 0, 0, 0], + dtype=torch.uint32, + device=device, + ) + routed_multicast_ptr = self._routed_symm_mem.multicast_ptr + if routed_multicast_ptr is None or routed_multicast_ptr == 0: + raise RuntimeError("routed NVLS multicast mapping is unavailable") + self._routed_multicast_ptr = int(routed_multicast_ptr) + + self._latent_output = torch.empty( + (max_m, latent_dim), dtype=torch.bfloat16, device=device + ) + self._shared_output = torch.empty( + (max_m, hidden_dim), dtype=torch.bfloat16, device=device + ) + shard_start = rank * self.shard_dim + shard_end = shard_start + self.shard_dim + self._shared_shard = self._shared_output[:, shard_start:shard_end] + + self._shared_workspace = symm_mem.empty( + (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim), + dtype=torch.bfloat16, + device=device, + ) + self._shared_symm_mem = symm_mem.rendezvous(self._shared_workspace, group) + self._shared_workspace.view(torch.int32).fill_(-0x80000000) + self._shared_flags = torch.zeros(12, dtype=torch.int32, device=device) + self._shared_flags[1] = 1 + self._shared_flags[2] = max_m * tp_size * self.shard_dim * 2 + peer_ptrs = [ + self._shared_symm_mem.get_buffer( + peer, + self._shared_workspace.shape, + torch.bfloat16, + ).data_ptr() + for peer in range(tp_size) + ] + if any(pointer == 0 for pointer in peer_ptrs): + raise RuntimeError("shared LSA peer mapping is unavailable") + self._shared_peer_ptrs = torch.tensor( + peer_ptrs, dtype=torch.int64, device=device + ) + + torch.accelerator.synchronize(device) + dist.barrier(group=group, device_ids=[device.index]) + for owner in range(tp_size): + if rank == owner: + compile_kernel( + rank=rank, + tp_size=tp_size, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + max_m=max_m, + max_token_ctas=max_token_ctas, + latent_output=self._latent_output, + routed_workspace=self._routed_workspace, + routed_flags=self._routed_flags, + routed_multicast_ptr=self._routed_multicast_ptr, + shared_output=self._shared_output, + shared_workspace=self._shared_workspace, + shared_flags=self._shared_flags, + shared_peer_ptrs=self._shared_peer_ptrs, + rms_eps=self.rms_eps, + fp32_internal=fp32_internal, + ) + dist.barrier(group=group, device_ids=[device.index]) + + def __call__( + self, + latent_source: torch.Tensor, + shared_source: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if latent_source.ndim != 2 or shared_source.ndim != 2: + raise ValueError("latent_source and shared_source must be rank-2") + m = latent_source.shape[0] + device = self._routed_workspace.device + expected = ( + (latent_source, (m, self.latent_dim), "latent_source"), + (shared_source, (m, self.hidden_dim), "shared_source"), + (gamma, (self.latent_dim,), "gamma"), + ) + for tensor, shape, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != torch.bfloat16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be contiguous CUDA BF16 {list(shape)}") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + with torch.accelerator.device_index(device.index): + launch( + latent_source, + gamma, + self._latent_output, + self._routed_workspace, + self._routed_flags, + self._routed_multicast_ptr, + shared_source, + self._shared_output, + self._shared_workspace, + self._shared_flags, + self._shared_peer_ptrs, + self.rms_eps, + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + max_token_ctas=self.max_token_ctas, + fp32_internal=self.fp32_internal, + ) + return ( + self._latent_output[:m], + self._shared_shard, + ) + + @property + def latent_output(self) -> torch.Tensor: + return self._latent_output + + @property + def shared_output(self) -> torch.Tensor: + return self._shared_output diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py new file mode 100644 index 00000000000..6a1ba9d77b4 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -0,0 +1,1291 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +"""Blackwell GEMM with a fused shared-shard add and multicast epilogue. + +Modified from the original CUTLASS CuTe DSL SM100 persistent GEMM tutorial. +The PDL wait before A loading orders the up-projection and shared-expert +inputs after the producer collective. +""" + +import math +from typing import Any + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.common import CacheEvictionPriority +from cutlass.cute.runtime import from_dlpack +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .fused_add_multicast_skinny_gemm import ( + FusedAddMulticastSkinnyGemmKernel, +) +from .primitives import CUDAGraphCompatibleWrapper + + +def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False): + converted = from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=16 + ) + if dynamic_m: + converted = converted.mark_compact_shape_dynamic( + mode=1, + stride_order=tensor.dim_order(), + ) + return converted + + +def validate_configuration( + *, + latent_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, +) -> None: + """Validate constraints imposed by this BF16 SM100 GEMM.""" + + if latent_dim <= 0 or shard_dim <= 0: + raise ValueError("GEMM K and N must be positive") + if latent_dim % 8 or shard_dim % 8: + raise ValueError("GEMM K and N must be divisible by 8 BF16 values") + if mma_tiler_mn[0] not in (64, 128): + raise ValueError("MMA M tile must be 64 or 128") + if mma_tiler_mn[1] not in range(32, 257, 32): + raise ValueError("MMA N tile must be a multiple of 32 in [32, 256]") + if ( + len(cluster_shape_mn) != 2 + or any(value <= 0 or value & (value - 1) for value in cluster_shape_mn) + or math.prod(cluster_shape_mn) > 16 + ): + raise ValueError("cluster dimensions must be powers of two with product <= 16") + if not 0 <= b_prime_stages <= math.ceil(latent_dim / 128): + raise ValueError("b_prime_stages exceeds the GEMM K-tile count") + + +@cute.jit +def _epilogue_tma_store_add_shared( + gemm_kernel, + epi_tidx: cutlass.Int32, + warp_idx: cutlass.Int32, + tma_atom_c: cute.CopyAtom, + tCtAcc_base: cute.Tensor, + sC: cute.Tensor, + tCgC_base: cute.Tensor, + tCgShared_base: cute.Tensor, + tCcC_base: cute.Tensor, + mC_mnl: cute.Tensor, + mShared_mnl: cute.Tensor, + epi_tile: cute.Tile, + num_tiles_executed: cutlass.Int32, + mma_tile_coord_mnl, + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + c_pipeline: pipeline.PipelineTmaStore, +) -> pipeline.PipelineState: + """BF16 GEMM rounding + BF16 shared shard addition, then swizzled S2G TMA.""" + sm100 = utils.gemm.sm100 + tCgC = sm100.transform_partitioned_tensor_layout(tCgC_base) + tCgShared = sm100.transform_partitioned_tensor_layout(tCgShared_base) + tCcC = sm100.transform_partitioned_tensor_layout(tCcC_base) + tCtAcc = sm100.transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = sm100.epilogue_tmem_copy_and_partition( + gemm_kernel, + epi_tidx, + tCtAcc, + tCgC, + epi_tile, + False, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tTR_rShared = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = sm100.epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gShared_partitioned = thr_copy_t2r.partition_D( + cute.flat_divide(tCgShared, epi_tile) + ) + tTR_cC_partitioned = thr_copy_t2r.partition_D(cute.flat_divide(tCcC, epi_tile)) + tTR_gShared = tTR_gShared_partitioned[ + (None, None, None, None, None, *mma_tile_coord_mnl) + ] + tTR_cC = tTR_cC_partitioned[(None, None, None, None, None, *mma_tile_coord_mnl)] + + exemplar = tTR_gShared_partitioned[(None, None, None, 0, 0, 0, 0, 0)] + mcl_r = cute.max_common_layout(tTR_rShared.layout, exemplar.layout) + shared_copy_bits = min( + exemplar.iterator.alignment * 8, + cute.size(mcl_r) * gemm_kernel.c_dtype.width, + 128, + ) + shared_g2r_atom = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + gemm_kernel.c_dtype, + num_bits_per_copy=shared_copy_bits, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gShared = cute.group_modes(tTR_gShared, 3, cute.rank(tTR_gShared)) + tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + previous_subtile_count = num_tiles_executed * subtile_cnt + cute.arch.griddepcontrol_wait() + for subtile_idx in range(subtile_cnt): + tTR_gShared_subtile = tTR_gShared[(None, None, None, subtile_idx)] + tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] + pred_shape = (1, *tTR_cC_subtile.shape[1:]) + pred = cute.make_rmem_tensor(pred_shape, cutlass.Boolean) + for m_idx in range(tTR_cC_subtile.shape[1]): + for n_idx in range(tTR_cC_subtile.shape[2]): + pred[(0, m_idx, n_idx)] = cute.elem_less( + tTR_cC_subtile[(0, m_idx, n_idx)], mC_mnl.shape + ) + tTR_rShared.store(cute.zeros_like(tTR_rShared, dtype=gemm_kernel.c_dtype)) + cute.copy( + shared_g2r_atom, + tTR_gShared_subtile, + tTR_rShared, + pred=pred, + ) + + # Load the shared addend before waiting for the accumulator. + if subtile_idx == 0: + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + gemm_vec = tiled_copy_r2s.retile(tTR_rAcc).load().to(gemm_kernel.c_dtype) + shared_vec = tiled_copy_r2s.retile(tTR_rShared).load() + fused_vec = (gemm_vec.to(cutlass.Float32) + shared_vec.to(cutlass.Float32)).to( + gemm_kernel.c_dtype + ) + # The symmetric output is an in-band Lamport mailbox whose empty + # marker contains BF16 -0. Normalize either signed zero to +0 so a + # legitimate result can never be mistaken for an unwritten fragment. + fused_vec = cute.where( + fused_vec == cute.zeros_like(fused_vec), + cute.zeros_like(fused_vec), + fused_vec, + ) + tRS_rC.store(fused_vec) + + c_buffer = (previous_subtile_count + subtile_idx) % gemm_kernel.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + epilog_sync_barrier.arrive_and_wait() + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: tuple[int, int, int], + a_dtype, + b_dtype, + c_dtype, + smem_capacity: int, + c_smem_layout, +) -> tuple[int, int, int]: + """Choose accumulator, mainloop, and epilogue stage counts.""" + num_acc_stage = 2 + num_c_stage = 2 + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + num_ab_stage = ( + smem_capacity - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + num_c_stage += ( + smem_capacity + - ab_bytes_per_stage * num_ab_stage + - (mbar_helpers_bytes + c_bytes) + ) // c_bytes_per_stage + return num_acc_stage, num_ab_stage, num_c_stage + + +class FusedAddMulticastGemm: + """Persistent Blackwell GEMM with a shared-add epilogue. + + B priming may overlap the producer collective. The PDL wait before A + loading orders both inputs before the shared-add epilogue. + """ + + def __init__( + self, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int = 2, + ): + self.acc_dtype = cutlass.Float32 + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + # B primes the combined A+B pipeline before the PDL wait. + self.b_prime_stages = b_prime_stages + self.cta_group = tcgen05.CtaGroup.ONE + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * len( + (self.mma_warp_id, self.tma_warp_id, *self.epilogue_warp_id) + ) + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + + def _create_tiled_mma(self): + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Derive layouts and stage counts from the compiled tensor shapes.""" + tiled_mma = self._create_tiled_mma() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + False, + self.c_layout, + self.c_dtype, + ) + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + utils.get_smem_capacity_in_bytes(), + c_smem_layout, + ) + + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, "sm_100" + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + c_multicast_i64: cutlass.Int64, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Launch the persistent GEMM.""" + # Preserve C's logical strided layout but point the TMA descriptor at + # this rank's shard inside the LSA multicast mapping. One TMA store is + # therefore replicated into the same shard on all eight ranks. + c = cute.make_tensor( + cute.make_ptr( + c.element_type, + c_multicast_i64, + cute.AddressSpace.gmem, + assumed_align=16, + ), + c.layout, + ) + + self.a_dtype: type[cutlass.Numeric] = a.element_type + self.b_dtype: type[cutlass.Numeric] = b.element_type + self.c_dtype: type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + tiled_mma = self._create_tiled_mma() + + self._setup_attributes() + if cutlass.const_expr(self.b_prime_stages > self.num_ab_stage): + raise ValueError( + "b_prime_stages exceeds the compiled A/B pipeline stage count" + ) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + tile_sched_params, + shared_shard, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=True, + ) + return + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + self._gemm_device( + tiled_mma, + tma_atom_a, + mA_mkl, + tma_atom_b, + mB_nkl, + tma_atom_c, + mC_mnl, + cluster_layout_vmnk, + a_smem_layout_staged, + b_smem_layout_staged, + c_smem_layout_staged, + epi_tile, + tile_sched_params, + shared_shard, + ) + + @cute.jit + def _gemm_device( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.Layout | cute.ComposedLayout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + shared_shard: cute.Tensor, + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + ab_producer, ab_consumer = ab_pipeline.make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + # Shared shard is physically [M, shard_dim]. Give it the same logical MNL + # view as C so its epilogue partition is coordinate-identical. + mShared_mnl = cute.make_tensor( + shared_shard.iterator, + cute.append(shared_shard.layout, cute.make_layout((1,), stride=(0,))), + ) + gShared_mnl = cute.local_tile( + mShared_mnl, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgC = thr_mma.partition_C(gC_mnl) + tCgShared = thr_mma.partition_C(gShared_mnl) + + # Predicate the partial M tile when M is smaller than the MMA tile. + idC = cute.make_identity_tensor(mC_mnl.shape) + cC_mnl = cute.local_tile( + idC, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + tCcC = thr_mma.partition_C(cC_mnl) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + gemm_grid_z = cute.arch.grid_dim()[2] + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + ( + cute.arch.grid_dim()[0], + cute.arch.grid_dim()[1], + gemm_grid_z, + ), + ) + work_tile = tile_sched.initial_work_tile_info() + + if warp_idx == self.tma_warp_id: + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Prime a short prefix of the existing combined A+B ring with + # B only. Its barrier still expects A+B bytes, so the MMA + # consumer cannot observe a half-filled stage. + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < self.b_prime_stages: + peek_ab_empty_status = ab_producer.try_acquire() + + cute.arch.griddepcontrol_wait() + + # Supply A to the very same stages after the producer AR has + # programmatically released this dependent kernel. + a_fill_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + for k_tile in cutlass.range(0, self.b_prime_stages, 1, unroll=1): + a_barrier = ab_pipeline.producer_get_barrier(a_fill_state) + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_tile)], + tAsA[(None, a_fill_state.index)], + tma_bar_ptr=a_barrier, + mcast_mask=a_full_mcast_mask, + ) + a_fill_state.advance() + + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range( + self.b_prime_stages, k_tile_cnt, 1, unroll=1 + ): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + ab_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in range(k_tile_cnt): + if is_leader_cta: + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + handle.release() + + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + acc_pipeline.producer_tail(acc_producer_state) + + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, producer_group=c_producer_group + ) + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + num_tiles_executed = tile_sched.num_tiles_executed + acc_consumer_state = _epilogue_tma_store_add_shared( + self, + tidx, + warp_idx, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + tCgShared, + tCcC, + mC_mnl, + mShared_mnl, + epi_tile, + num_tiles_executed, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + c_pipeline, + ) + + c_pipeline.producer_tail() + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # Allow the Lamport copy to become resident before this grid + # fully retires. Its griddepcontrol.wait still enforces complete + # producer-grid ordering before mailbox inspection. + cute.arch.griddepcontrol_launch_dependents() + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + max_active_clusters: cutlass.Constexpr, + ) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]: + """Build the static persistent schedule.""" + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: tuple[int, int, int], + num_acc_stage: int, + arch: str, + ) -> int: + """Return the required tensor-memory column count.""" + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) + + return num_tmem_alloc_cols + + +@cute.jit +def launch_kernel( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, n, k) + c: cute.Tensor, # (l, m, n) + shared_shard: cute.Tensor, # (m, shard_dim), private TP shard + rows: cutlass.Int64, + c_multicast_i64: cutlass.Int64, + full_hidden_dim: cutlass.Constexpr, + shard_dim: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, +): + """Launch the fused-add multicast GEMM using PyTorch BMM tensor order.""" + # C is passed as the fixed-capacity [1,max_m,H] symmetric mailbox. Only M + # is runtime-variable; construct this rank's logical [1,M,S] view here so + # H and S remain compile-time constants and no dynamic strided host view is + # needed on every call. __call__ later replaces the local base pointer with + # the already rank-offset multicast address. + c = cute.make_tensor( + c.iterator, + cute.make_layout( + (1, rows, shard_dim), + stride=(0, full_hidden_dim, 1), + ), + ) + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,n,k) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[1, 2, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op( + a, + b, + c, + shared_shard, + c_multicast_i64, + max_active_clusters, + stream, + ) + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + mnkl: tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + shared_shard: cute.Tensor, + full_hidden_dim: int, + shard_dim: int, + mma_tiler_mn: tuple[int, int] = (64, 32), + cluster_shape_mn: tuple[int, int] = (1, 8), + max_active_clusters: cutlass.Constexpr = None, + b_prime_stages: int = 2, +): + key = ( + torch.accelerator.current_device_index(), + mnkl, + full_hidden_dim, + shard_dim, + mma_tiler_mn, + cluster_shape_mn, + max_active_clusters, + b_prime_stages, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + gemm = FusedAddMulticastGemm( + mma_tiler_mn, + cluster_shape_mn, + b_prime_stages, + ) + validate_configuration( + latent_dim=mnkl[2], + shard_dim=mnkl[1], + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if any(tensor.element_type is not cutlass.BFloat16 for tensor in (a, b, c)): + raise ValueError("up-projection tensors must be BF16") + + # The producer writes [M, full_hidden_dim]. GEMM receives a rank-local + # [M, shard_dim] view: contiguous within a row, with full_hidden_dim as + # its leading dimension. + shared_shard_compile = make_fake_tensor( + shared_shard.element_type, + (mnkl[0], shard_dim), + stride=(full_hidden_dim, 1), + assumed_align=16, + ) + stream = make_fake_stream() + compiled = cute.compile( + launch_kernel, + gemm, + a, + b, + c, + shared_shard_compile, + cutlass.Int64(mnkl[0]), + cutlass.Int64(0), + full_hidden_dim, + shard_dim, + max_active_clusters, + stream, + ) + _COMPILED[key] = compiled + return compiled + + +class AdaptiveUpProjectionKernel: + """Dispatch static-M Skinny or dynamic-M WGMMA into one mailbox.""" + + def __init__( + self, + *, + group: dist.ProcessGroup, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + max_m: int, + skinny_max_m: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + b_prime_stages: int, + ) -> None: + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + if not 0 <= skinny_max_m <= min(8, max_m): + raise ValueError("skinny_max_m must be in [0, min(8, max_m)]") + self.rank = rank + self.tp_size = tp_size + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.max_m = max_m + self.skinny_max_m = skinny_max_m + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + self.b_prime_stages = b_prime_stages + device = torch.device("cuda", torch.accelerator.current_device_index()) + self._device = device + self._dynamic: Any | None = None + self._skinny_by_m: dict[int, FusedAddMulticastSkinnyGemmKernel] = {} + validate_configuration( + latent_dim=latent_dim, + shard_dim=self.shard_dim, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + b_prime_stages=b_prime_stages, + ) + if skinny_max_m and self.latent_dim % (224 * 8): + raise ValueError( + "Skinny up-projection requires latent_dim divisible by 1792." + ) + + self._mailbox = symm_mem.empty( + (1, max_m, hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + self._mailbox_symm_mem = symm_mem.rendezvous(self._mailbox, group) + self._mailbox.view(torch.int32).fill_(-0x80000000) + multicast_ptr = self._mailbox_symm_mem.multicast_ptr + if multicast_ptr is None or multicast_ptr == 0: + raise RuntimeError("mailbox NVLS multicast mapping is unavailable") + self._mailbox_multicast_ptr = int(multicast_ptr) + + cluster_size = math.prod(cluster_shape_mn) + self._max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_size + ) + self._mailbox_c = _as_cute(self._mailbox) + + def compile_dynamic(self) -> None: + if self._dynamic is not None: + return + device = self._device + with torch.accelerator.device_index(device.index): + compile_latent = torch.empty( + (1, self.max_m, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_weight = torch.empty( + (self.shard_dim, self.latent_dim), + dtype=torch.bfloat16, + device=device, + ) + compile_shared = torch.empty( + (self.max_m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + )[ + :, + self.rank * self.shard_dim : (self.rank + 1) * self.shard_dim, + ] + compile_latent_c = _as_cute(compile_latent, dynamic_m=True) + compile_weight_c = _as_cute(compile_weight.unsqueeze(0)) + compile_shared_c = _as_cute(compile_shared) + + self._dynamic = compile_kernel( + (self.max_m, self.shard_dim, self.latent_dim, 1), + compile_latent_c, + compile_weight_c, + self._mailbox_c, + compile_shared_c, + self.hidden_dim, + self.shard_dim, + self.mma_tiler_mn, + self.cluster_shape_mn, + self._max_active_clusters, + self.b_prime_stages, + ) + + def compile_skinny(self, m: int) -> None: + if not 1 <= m <= self.skinny_max_m: + raise ValueError( + f"Skinny up-projection requires M in [1, {self.skinny_max_m}]." + ) + if m in self._skinny_by_m: + return + with torch.accelerator.device_index(self._device.index): + self._skinny_by_m[m] = FusedAddMulticastSkinnyGemmKernel( + rank=self.rank, + tp_size=self.tp_size, + latent_dim=self.latent_dim, + hidden_dim=self.hidden_dim, + num_rows=m, + ) + + def ensure_compiled(self, m: int) -> None: + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + if m <= self.skinny_max_m: + self.compile_skinny(m) + else: + self.compile_dynamic() + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + ) -> torch.Tensor: + if latent.ndim != 2: + raise ValueError("latent must be rank-2") + m = latent.shape[0] + device = self._mailbox.device + expected = ( + (latent, (m, self.latent_dim), "latent"), + ( + weight, + (self.shard_dim, self.latent_dim), + "weight", + ), + ( + shared_shard, + (self.max_m, self.shard_dim), + "shared_shard", + ), + ) + for tensor, shape, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != torch.bfloat16 + or tensor.device != device + ): + raise ValueError(f"{name} must be CUDA torch.bfloat16 {list(shape)}") + if ( + not latent.is_contiguous() + or not weight.is_contiguous() + or shared_shard.stride() != (self.hidden_dim, 1) + ): + raise ValueError("up-projection inputs have unsupported strides") + if not 1 <= m <= self.max_m: + raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]") + + if m <= self.skinny_max_m: + skinny = self._skinny_by_m.get(m) + if skinny is None: + raise RuntimeError( + f"Skinny up-projection M={m} was not compiled before launch." + ) + return skinny( + latent, + weight, + shared_shard, + self._mailbox, + self._mailbox_multicast_ptr, + ) + + if self._dynamic is None: + raise RuntimeError("Dynamic up-projection was not compiled before launch.") + with torch.accelerator.device_index(device.index): + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + self._dynamic( + _as_cute(latent.unsqueeze(0), dynamic_m=True), + _as_cute(weight.unsqueeze(0)), + self._mailbox_c, + _as_cute(shared_shard), + cutlass.Int64(m), + cutlass.Int64( + self._mailbox_multicast_ptr + self.rank * self.shard_dim * 2 + ), + stream, + ) + return self._mailbox diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py new file mode 100644 index 00000000000..58d8f1ca18b --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Static-M SIMT skinny GEMM with a shared-add multicast epilogue.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int64, const_expr +from cutlass.cute.runtime import from_dlpack + +from .primitives import ( + CUDAGraphCompatibleWrapper, + bf16x2_to_u32, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + sanitize_negative_zero, + sanitize_negative_zero_u32, + sanitize_negative_zero_u32x2, + store_global_u32, + store_global_u32x2, + store_global_u32x4, +) + + +@dataclass(frozen=True) +class SkinnyConfig: + block_size: int = 224 + outputs_per_block: int = 2 + k_unroll: int = 2 + vector_width: int = 8 + prefetch_b_before_pdl: bool = True + + +def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig: + if shard_dim == 448: + if num_rows >= 6: + return SkinnyConfig( + block_size=224, + outputs_per_block=2, + k_unroll=1, + ) + outputs_per_block = 2 if num_rows <= 3 else 4 + return SkinnyConfig( + block_size=448, + outputs_per_block=outputs_per_block, + k_unroll=1, + ) + if num_rows == 1: + return SkinnyConfig(outputs_per_block=8, k_unroll=1) + return SkinnyConfig(outputs_per_block=4) + + +def _as_cute(tensor: torch.Tensor): + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), + assumed_align=16, + ) + + +class FusedAddMulticastSkinnyGemm: + """SIMT GEMM adapted from the existing Skinny GEMM.""" + + def __init__( + self, + *, + num_rows: int, + hidden_dim: int, + config: SkinnyConfig, + ) -> None: + if config.block_size % 32: + raise ValueError("skinny block_size must be a multiple of 32") + self.num_rows = num_rows + self.hidden_dim = hidden_dim + self.block_size = config.block_size + self.outputs_per_block = config.outputs_per_block + self.k_unroll = config.k_unroll + self.vector_width = config.vector_width + self.prefetch_b_before_pdl = config.prefetch_b_before_pdl + self.num_warps = config.block_size // 32 + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + stream: cuda.CUstream, + ) -> None: + n = cute.size(gB, mode=[0]) + k = cute.size(gA, mode=[1]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + BFloat16, + num_bits_per_copy=self.vector_width * BFloat16.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + gA, + gB, + gShared, + output_multicast_ptr, + k, + copy_a, + copy_b, + ).launch( + grid=[cute.ceil_div(n, self.outputs_per_block), 1, 1], + block=[self.block_size, 1, 1], + smem=(self.num_rows * self.outputs_per_block * self.num_warps * 4), + stream=stream, + use_pdl=True, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gShared: cute.Tensor, + output_multicast_ptr: Int64, + k_extent: cutlass.Int32, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + num_rows: cutlass.Constexpr = self.num_rows + + acc = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, outputs_per_block), + stride=(outputs_per_block, 1), + ), + Float32, + ) + acc.fill(0.0) + + n_base = block_idx * outputs_per_block + k_tile_size: cutlass.Constexpr = block_size * vector_width + num_k_tiles = k_extent // k_tile_size + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + a_regs = cute.make_rmem_tensor( + cute.make_layout( + (num_rows, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, vector_width), + stride=(vector_width, 1), + ), + BFloat16, + ) + + if const_expr(self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + + cute.arch.griddepcontrol_wait() + + for mi in cutlass.range_constexpr(num_rows): + cute.copy(copy_a, tA[mi, None, 0], a_regs[mi, None]) + if const_expr(not self.prefetch_b_before_pdl): + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, 0], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ + ni, vi + ].to(Float32) + + for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll): + for mi in cutlass.range_constexpr(num_rows): + cute.copy( + copy_a, + tA[mi, None, k_tile], + a_regs[mi, None], + ) + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, None]) + for vi in cutlass.range_constexpr(vector_width): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[ + ni, vi + ].to(Float32) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor( + Float32, + smem_layout, + byte_alignment=16, + ) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + fused = cute.make_rmem_tensor( + cute.make_layout((outputs_per_block,)), + BFloat16, + ) + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + total = ( + partials[mi, ni, None] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + ) + gemm_value = Float32(total).to(BFloat16) + fused[ni] = ( + gemm_value.to(Float32) + gShared[mi, n_base + ni].to(Float32) + ).to(BFloat16) + output_offset = Int64((mi * self.hidden_dim + n_base) * 2) + if const_expr(outputs_per_block == 2): + packed = sanitize_negative_zero_u32(bf16x2_to_u32(fused.load())) + store_global_u32( + output_multicast_ptr + output_offset, + packed, + ) + elif const_expr(outputs_per_block == 4): + packed = sanitize_negative_zero_u32x2( + bf16x4_to_packed_u32x2(fused.load()) + ) + store_global_u32x2( + output_multicast_ptr + output_offset, + packed, + ) + else: + packed = sanitize_negative_zero( + bf16x8_to_packed_u32x4(fused.load()) + ) + store_global_u32x4( + output_multicast_ptr + output_offset, + packed, + ) + + cute.arch.griddepcontrol_launch_dependents() + + +_COMPILED: dict[tuple[object, ...], object] = {} + + +def compile_kernel( + *, + num_rows: int, + latent_dim: int, + hidden_dim: int, + shard_dim: int, + config: SkinnyConfig, +): + key = ( + torch.accelerator.current_device_index(), + num_rows, + latent_dim, + hidden_dim, + shard_dim, + config, + ) + if key in _COMPILED: + return _COMPILED[key] + + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + if shard_dim % config.outputs_per_block: + raise ValueError("shard_dim must be divisible by outputs_per_block") + if latent_dim % (config.block_size * config.vector_width): + raise ValueError("latent_dim must be divisible by block_size * vector_width") + + a = make_fake_tensor( + BFloat16, + (num_rows, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + b = make_fake_tensor( + BFloat16, + (shard_dim, latent_dim), + stride=(latent_dim, 1), + assumed_align=16, + ) + shared = make_fake_tensor( + BFloat16, + (num_rows, shard_dim), + stride=(hidden_dim, 1), + assumed_align=16, + ) + compiled = cute.compile( + FusedAddMulticastSkinnyGemm( + num_rows=num_rows, + hidden_dim=hidden_dim, + config=config, + ), + a, + b, + shared, + Int64(0), + make_fake_stream(), + options="--ptxas-options -maxrregcount=128", + ) + _COMPILED[key] = compiled + return compiled + + +class FusedAddMulticastSkinnyGemmKernel: + """Buffer-free compiled launcher for one static M.""" + + def __init__( + self, + *, + rank: int, + tp_size: int, + latent_dim: int, + hidden_dim: int, + num_rows: int, + ) -> None: + if not 1 <= num_rows <= 8: + raise ValueError("skinny backend requires static M in [1, 8]") + if hidden_dim % tp_size: + raise ValueError("hidden_dim must be divisible by TP size") + self.rank = rank + self.latent_dim = latent_dim + self.hidden_dim = hidden_dim + self.shard_dim = hidden_dim // tp_size + self.num_rows = num_rows + self._skinny = compile_kernel( + num_rows=num_rows, + latent_dim=latent_dim, + hidden_dim=hidden_dim, + shard_dim=self.shard_dim, + config=config_for_m(num_rows, self.shard_dim), + ) + + def __call__( + self, + latent: torch.Tensor, + weight: torch.Tensor, + shared_shard: torch.Tensor, + mailbox: torch.Tensor, + mailbox_multicast_ptr: int, + ) -> torch.Tensor: + device = mailbox.device + expected = ( + ( + latent, + (self.num_rows, self.latent_dim), + torch.bfloat16, + "latent", + ), + ( + weight, + (self.shard_dim, self.latent_dim), + torch.bfloat16, + "weight", + ), + ( + shared_shard, + (mailbox.shape[1], self.shard_dim), + torch.bfloat16, + "shared_shard", + ), + ( + mailbox, + (1, mailbox.shape[1], self.hidden_dim), + torch.bfloat16, + "mailbox", + ), + ) + for tensor, shape, dtype, name in expected: + if ( + tensor.shape != shape + or tensor.dtype != dtype + or tensor.device != device + ): + raise ValueError(f"{name} must be CUDA {dtype} {list(shape)}") + if ( + not latent.is_contiguous() + or not weight.is_contiguous() + or shared_shard.stride() != (self.hidden_dim, 1) + or not mailbox.is_contiguous() + ): + raise ValueError("skinny up-projection inputs have unsupported strides") + if mailbox.shape[1] < self.num_rows: + raise ValueError("mailbox capacity is smaller than runtime M") + + with torch.accelerator.device_index(device.index): + self._skinny( + _as_cute(latent), + _as_cute(weight), + _as_cute(shared_shard[: self.num_rows]), + Int64(mailbox_multicast_ptr + self.rank * self.shard_dim * 2), + cuda.CUstream(torch.cuda.current_stream(device).cuda_stream), + ) + return mailbox diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py new file mode 100644 index 00000000000..e178e88c413 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Lamport mailbox-to-local copy and reset.""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from .primitives import ( + VEC_BF16, + fragment_is_dirty, + load_global_u32x4, + store_global_u32x4, + store_lamport_sentinel_128, + to_cute, + to_cute_dynamic_m, +) + + +class LamportCopy: + """Consume the local physical copy of an NVLS-multicast mailbox.""" + + def __init__(self, hidden_dim: int, ctas: int, threads: int): + self.hidden_dim = hidden_dim + self.ctas = ctas + self.threads = threads + + @cute.jit + def __call__( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + stream: cuda.CUstream, + ): + self.kernel(symmetric_mailbox, local_output, m).launch( + grid=(self.ctas, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + symmetric_mailbox: cute.Tensor, + local_output: cute.Tensor, + m: cutlass.Int32, + ): + # The CTA may be scheduled early, but mailbox inspection must not pass + # the producer GEMM's programmatic completion point. + cute.arch.griddepcontrol_wait() + + tidx, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + thread = cutlass.Int64(block * self.threads + tidx) + stride = cutlass.Int64(self.ctas * self.threads) + fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16) + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source, volatile=True) + while fragment_is_dirty(packed): + packed = load_global_u32x4(source, volatile=True) + + destination = cutlass.Int64((local_output.iterator + element).toint()) + store_global_u32x4(destination, packed, volatile=False) + fragment = fragment + stride + + # The returned ordinary tensor is complete. A same-stream successor + # may overlap the mailbox cleanup below. + cute.arch.griddepcontrol_launch_dependents() + + fragment = thread + while fragment < fragments: + element = fragment * VEC_BF16 + source = cute.make_ptr( + cutlass.BFloat16, + (symmetric_mailbox.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + store_lamport_sentinel_128(source) + fragment = fragment + stride + + +@functools.cache +def compile_kernel( + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + device_index: int, +): + if hidden_dim <= 0 or hidden_dim % VEC_BF16: + raise ValueError("hidden_dim must be a positive multiple of 8") + if max_m <= 0: + raise ValueError("max_m must be positive") + if ctas <= 0: + raise ValueError("Lamport copy ctas must be positive") + if not 1 <= threads <= 1024: + raise ValueError("Lamport copy threads must be in [1, 1024]") + with torch.accelerator.device_index(device_index): + mailbox = make_fake_compact_tensor( + cutlass.BFloat16, (max_m * hidden_dim,), assumed_align=16 + ) + output = make_fake_compact_tensor( + cutlass.BFloat16, + (cute.sym_int32(divisibility=VEC_BF16),), + assumed_align=16, + ) + return cute.compile( + LamportCopy(hidden_dim, ctas, threads), + mailbox, + output, + cutlass.Int32(max_m), + make_fake_stream(), + ) + + +def launch( + symmetric_mailbox: torch.Tensor, + local_output: torch.Tensor, + *, + m: int, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, +) -> None: + if not 1 <= m <= max_m: + raise ValueError(f"runtime M={m} must be in [1, {max_m}]") + if ( + symmetric_mailbox.ndim != 3 + or symmetric_mailbox.dtype != torch.bfloat16 + or not symmetric_mailbox.is_cuda + or not symmetric_mailbox.is_contiguous() + or symmetric_mailbox.shape[0] != 1 + or symmetric_mailbox.shape[2] != hidden_dim + ): + raise ValueError( + f"symmetric_mailbox must be contiguous CUDA BF16 [1,M,{hidden_dim}]" + ) + if symmetric_mailbox.shape[1] != max_m: + raise ValueError("symmetric mailbox must retain its full max_m capacity") + if ( + local_output.shape != (1, m, hidden_dim) + or local_output.dtype != torch.bfloat16 + or local_output.device != symmetric_mailbox.device + or not local_output.is_contiguous() + ): + raise ValueError("local_output must be contiguous CUDA BF16 [1,M,H]") + + device_index = symmetric_mailbox.device.index + stream = cuda.CUstream( + torch.cuda.current_stream(symmetric_mailbox.device).cuda_stream + ) + compile_kernel(hidden_dim, max_m, ctas, threads, device_index)( + to_cute(symmetric_mailbox.flatten(), 16), + to_cute_dynamic_m( + local_output.flatten(), + mode=0, + assumed_align=16, + ), + cutlass.Int32(m), + stream, + ) + + +class LamportCopyKernel: + """Copy a borrowed symmetric mailbox into a fresh local tensor.""" + + def __init__( + self, + *, + hidden_dim: int, + max_m: int, + ctas: int, + threads: int, + ) -> None: + self.hidden_dim = hidden_dim + self.max_m = max_m + self.ctas = ctas + self.threads = threads + compile_kernel( + hidden_dim, + max_m, + ctas, + threads, + torch.accelerator.current_device_index(), + ) + + def __call__(self, symmetric_mailbox: torch.Tensor, *, m: int) -> torch.Tensor: + if not symmetric_mailbox.is_cuda: + raise ValueError("symmetric_mailbox must be a CUDA tensor") + device = symmetric_mailbox.device + with torch.accelerator.device_index(device.index): + output = torch.empty( + (1, m, self.hidden_dim), + dtype=torch.bfloat16, + device=device, + ) + launch( + symmetric_mailbox, + output, + m=m, + hidden_dim=self.hidden_dim, + max_m=self.max_m, + ctas=self.ctas, + threads=self.threads, + ) + return output diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py new file mode 100644 index 00000000000..1bcb6d400cc --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared CuTe DSL primitives; this module does not define a CUDA kernel.""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +VEC_BF16 = 8 +PACKED_BYTES = 16 +NUM_LAMPORT_BUFFERS = 3 +NEG_ZERO_F32_BITS = 0x80000000 +NEG_ZERO_BF16_BITS = 0x8000 + + +class CUDAGraphCompatibleWrapper: + """DLPack view that does not synchronize with the producer stream.""" + + def __init__(self, tensor: torch.Tensor): + self.tensor = tensor + + def __dlpack__(self, stream=None): + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor: + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align + ) + + +def to_cute_dynamic_m( + tensor: torch.Tensor, + *, + mode: int, + assumed_align: int = 16, +) -> cute.Tensor: + """Expose exactly one compact tensor mode as a runtime shape. + + Model dimensions remain part of the compiled tensor type. Only the token + mode is symbolic, so changing M within an Op's capacity reuses the same + compiled kernel. + """ + + return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( + mode=mode, + stride_order=tensor.dim_order(), + ) + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + """Load one 128-bit fragment as four u32 registers. + + The volatile form is the Lamport polling load. Marking the asm + side-effecting prevents loop-invariant motion and common-subexpression + elimination across polling iterations. + """ + + address = pointer.toint(loc=loc, ip=ip) + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=volatile, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def store_global_u32x4( + address: Int64, + packed, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +) -> None: + """Store four packed words to an ordinary or NVLS multicast global VA.""" + + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] + opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + f"{opcode} [$0], {{$1, $2, $3, $4}};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32x2( + address: Int64, + packed, + *, + loc=None, + ip=None, +) -> None: + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32( + address: Int64, + word: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + word.ir_value(loc=loc, ip=ip), + ], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: + """Reset one Lamport fragment to four FP32 negative-zero bit patterns.""" + + address = pointer.toint(loc=loc, ip=ip) + value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value, value, value, value], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_release_gpu_add_u32( + pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None +) -> None: + """The exact SM100 arrival primitive used by upstream LamportFlags.""" + + address = pointer.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "red.async.release.global.gpu.add.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def map_shared_to_peer( + smem_ptr: cute.Pointer, + peer_rank: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Map a local shared-memory slot to the same slot in a peer CTA.""" + + smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return Int32( + llvm.inline_asm( + T.i32(), + [smem_address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, + value: Float32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + remote_address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc) + values = llvm.bitcast( + target_type, + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, BFloat16) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([4], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + target_type = ir.VectorType.get([2], T.i32(), loc=loc) + packed = llvm.bitcast( + target_type, + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def bf16x2_to_u32(values, *, loc=None, ip=None): + packed = llvm.bitcast( + T.i32(), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return Uint32(packed) + + +@cute.jit +def sanitize_negative_zero_u32(word): + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for i in cutlass.range_constexpr(2): + result[i] = sanitize_negative_zero_u32(packed[i]) + return result.load() + + +@cute.jit +def sanitize_negative_zero(packed): + """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" + + result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for i in cutlass.range_constexpr(4): + word = packed[i] + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEG_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + result[i] = word + return result.load() + + +@cute.jit +def fragment_is_dirty(packed): + """Bit-exact upstream sentinel check: one comparison per 32-bit word.""" + + dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS) + for i in cutlass.range_constexpr(1, 4): + dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS)) + return dirty + + +@cute.jit +def warp_sum_specialized( + value: Float32, + warp_idx: Int32, + lane: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Warp sum supporting a compile-time partial final warp.""" + + if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): + for offset in cutlass.range_constexpr(16, 0, -1): + # range_constexpr does not provide powers-of-two stepping. + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + other = cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=last_warp_mask, + mask_and_clamp=31, + ) + if (lane ^ Int32(offset)) < Int32(last_warp_lanes): + value = value + other + else: + for offset in cutlass.range_constexpr(16, 0, -1): + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + return value + + +@cute.jit +def block_sum_specialized( + value: Float32, + warp_sums: cute.Tensor, + tidx: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Upstream-equivalent FP32 block reduction.""" + + lane = cute.arch.lane_idx() + warp_idx = cute.arch.warp_idx() + value = warp_sum_specialized( + value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask + ) + if lane == 0: + warp_sums[warp_idx] = value + cute.arch.barrier() + + block_sum = Float32(0.0) + if warp_idx == 0: + if lane < Int32(warps): + block_sum = warp_sums[lane] + block_sum = cute.arch.warp_reduction_sum(block_sum) + if lane == 0: + warp_sums[0] = block_sum + cute.arch.barrier() + return warp_sums[0] diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py new file mode 100644 index 00000000000..816da37153f --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MLA prefill and decode epilogues for Kimi-K3. + +Thin wrappers over the CUDA ops in +``csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu``, which +mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. + +- ``fused_mla_key_concat_kv_cache_insert`` (bf16): optionally apply RoPE, + concat the full per-head key ``[k_nope | k_pe]`` into ``k_out``, and insert + the latent ``[kv_c_normed | k_pe]`` into the paged cache. +- ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize + ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the + cache shares ``k_scale``, as in ``concat_and_cache_mla``). + +The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE +inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use +Programmatic Dependent Launch to overlap the tail of the producing GEMMs on +sm_90+. +""" + +import torch + + +def fused_mla_key_concat_kv_cache_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Apply optional RoPE, concat K, and insert the paged latent (bf16). + + Returns the full key ``[Tp, H, qk_nope_head_dim + qk_rope_head_dim]``; + optionally rotates ``q`` and writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty( + (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device + ) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_kv_cache_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_key_concat_ds_mla_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim], RoPE is applied in place + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + kv_cache: torch.Tensor, # [num_blocks, block_size, 656] uint8 (fp8_ds_mla) + slot_mapping: torch.Tensor, # [Tp] int64 + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat full K (bf16) and insert the latent in the fp8_ds_mla layout. + + The cache uses DeepSeek's 656-byte block-scaled layout (NoPE in 4 tiles of + 128 with per-tile dynamic fp8 scales, RoPE as bf16) -- self-scaling, so no + scale argument. Returns the bf16 full key; optionally rotates ``q`` and + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, qk_nope_head_dim = k_nope.shape + qk_head_dim = qk_nope_head_dim + k_pe.shape[1] + k_out = torch.empty( + (tp, num_heads, qk_head_dim), dtype=k_nope.dtype, device=k_nope.device + ) + if tp == 0: + return k_out + torch.ops._C.fused_kimi_k3_mla_key_concat_ds_mla_insert( + q, + k_nope, + k_pe, + kv_c_normed, + k_out, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return k_out + + +def fused_mla_qkv_quant_kv_cache_fp8_insert( + q: torch.Tensor, # [Tp, H, qk_head_dim] + k_nope: torch.Tensor, # [Tp, H, qk_nope_head_dim] + k_pe: torch.Tensor, # [Tp, qk_rope_head_dim] or [Tp, 1, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [Tp, kv_lora_rank] + v: torch.Tensor, # [Tp, H, v_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, kv_lora_rank + rope] fp8 + slot_mapping: torch.Tensor, # [Tp] int64 + q_scale_inv: torch.Tensor, # scalar fp32, 1 / q scale (attention query) + k_scale_inv: torch.Tensor, # scalar fp32, 1 / k scale (attention key) + v_scale_inv: torch.Tensor, # scalar fp32, 1 / v scale (attention value) + cache_scale_inv: torch.Tensor, # scalar fp32, 1 / kv scale (cache latent) + positions: torch.Tensor | None = None, # [Tp] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize q/k/v to fp8 and insert the fp8 latent into the paged cache. + + The attention key ``k_fp8`` and the cache latent use *separate* scales + (``k_scale_inv`` vs ``cache_scale_inv``): the cache must be quantized with + ``_k_scale`` (read back by decode / context), while the prefill attention + q/k/v currently stay unscaled (the prefill flash path does not dequantize). + + Returns ``(q_fp8, k_fp8, v_fp8)``; writes the fp8 ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + tp, num_heads, _ = q.shape + qk_head_dim = q.shape[2] + v_head_dim = v.shape[2] + fp8 = torch.float8_e4m3fn + q_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + k_fp8 = torch.empty((tp, num_heads, qk_head_dim), dtype=fp8, device=q.device) + v_fp8 = torch.empty((tp, num_heads, v_head_dim), dtype=fp8, device=q.device) + if tp == 0: + return q_fp8, k_fp8, v_fp8 + torch.ops._C.fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( + q, + k_nope, + k_pe, + kv_c_normed, + v, + q_fp8, + k_fp8, + v_fp8, + kv_cache, + slot_mapping, + q_scale_inv, + k_scale_inv, + v_scale_inv, + cache_scale_inv, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return q_fp8, k_fp8, v_fp8 + + +def fused_mla_decode_q_concat_kv_cache_insert( + ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q) + q_pe: torch.Tensor, # [B, H, qk_rope_head_dim] + kv_c_normed: torch.Tensor, # [B, kv_lora_rank] + k_pe: torch.Tensor, # [B, qk_rope_head_dim] or [B, 1, qk_rope_head_dim] + kv_cache: torch.Tensor, # [num_blocks, block_size, entry] + slot_mapping: torch.Tensor, # [B] int64 + *, + ds_mla: bool = False, + q_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / q scale + cache_scale_inv: torch.Tensor | None = None, # scalar fp32, 1 / kv scale + positions: torch.Tensor | None = None, # [B] int64 + cos_sin_cache: torch.Tensor | None = None, # [max_position, rope] +) -> torch.Tensor: + """Concat the latent decode query ``mqa_q = [ql_nope | q_pe]`` and insert the + latent ``[kv_c_normed | k_pe]`` into the paged cache, in one launch (runs + right before ``forward_mqa``). + + Dispatched by cache format: + - bf16 -> bf16 mqa_q, bf16 cache + - plain fp8 -> fp8 mqa_q (q_scale_inv), fp8 cache (cache_scale_inv) + - fp8_ds_mla -> bf16 mqa_q, 656B block-scaled cache + + Returns ``mqa_q`` of shape ``[B, H, kv_lora_rank + qk_rope_head_dim]``; + writes ``kv_cache`` in place. + """ + k_pe = k_pe.reshape(k_pe.shape[0], -1) + b, num_heads, kv_lora_rank = ql_nope.shape + entry = kv_lora_rank + q_pe.shape[-1] + fp8_q = q_scale_inv is not None + out_dtype = torch.float8_e4m3fn if fp8_q else ql_nope.dtype + mqa_q = torch.empty((b, num_heads, entry), dtype=out_dtype, device=ql_nope.device) + if b == 0: + return mqa_q + + if ds_mla: + cache = ( + kv_cache if kv_cache.dtype == torch.uint8 else kv_cache.view(torch.uint8) + ) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_ds_mla_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + cache.shape[1], + positions, + cos_sin_cache, + ) + elif fp8_q: + assert cache_scale_inv is not None, "fp8 decode requires cache_scale_inv" + cache = ( + kv_cache + if kv_cache.dtype == torch.float8_e4m3fn + else kv_cache.view(torch.float8_e4m3fn) + ) + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + cache, + slot_mapping, + q_scale_inv, + cache_scale_inv, + cache.shape[1], + positions, + cos_sin_cache, + ) + else: + torch.ops._C.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( + ql_nope, + q_pe, + kv_c_normed, + k_pe, + mqa_q, + kv_cache, + slot_mapping, + kv_cache.shape[1], + positions, + cos_sin_cache, + ) + return mqa_q diff --git a/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py new file mode 100644 index 00000000000..fe97c8fd308 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from functools import partial +from typing import ClassVar + +import torch +import torch.distributed as dist + +from vllm.distributed import get_tp_group +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) + +_MAX_NUM_TOKENS = 16 +_SKINNY_MAX_NUM_TOKENS = 5 +_MMA_TILER_MN = (64, 32) +_GEMM_CLUSTER_MN = (1, 8) +_B_PRIME_STAGES = 2 +_COLLECTIVE_TOKEN_CTAS = 8 +_LAMPORT_COPY_CTAS = 32 +_LAMPORT_COPY_THREADS = 224 +_SUPPORTED_TP_SIZES = (8, 16) + + +@dataclass(frozen=True) +class KimiK3LatentMoETailContract: + tp_group_id: int + tp_size: int + device: torch.device + dtype: torch.dtype + hidden_size: int + latent_size: int + max_num_tokens: int + rms_eps: float + + +class KimiK3LatentMoETailOp: + """Process-wide cached K3 latent-MoE tail implementation.""" + + _instances: ClassVar[ + dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"] + ] = {} + + @classmethod + def _contract_and_group( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]: + tp = get_tp_group() + group = tp.device_group + device = torch.device(device) + tp_device = torch.device(tp.device) + if device != tp_device: + raise ValueError( + f"Input device {device} does not match TP device {tp_device}." + ) + return ( + KimiK3LatentMoETailContract( + tp_group_id=id(group), + tp_size=dist.get_world_size(group), + device=device, + dtype=dtype, + hidden_size=hidden_size, + latent_size=latent_size, + max_num_tokens=_MAX_NUM_TOKENS, + rms_eps=float(rms_eps), + ), + group, + ) + + @classmethod + def initialize( + cls, + *, + hidden_size: int, + latent_size: int, + dtype: torch.dtype, + device: torch.device, + rms_eps: float, + ) -> "KimiK3LatentMoETailOp": + contract, group = cls._contract_and_group( + hidden_size=hidden_size, + latent_size=latent_size, + dtype=dtype, + device=device, + rms_eps=rms_eps, + ) + op = cls._instances.get(contract) + if op is None: + op = cls(contract, group) + cls._instances[contract] = op + return op + + def __init__( + self, + contract: KimiK3LatentMoETailContract, + group: dist.ProcessGroup, + ) -> None: + if contract.tp_size not in _SUPPORTED_TP_SIZES: + raise ValueError( + "K3 latent-MoE tail fusion requires TP in " + f"{_SUPPORTED_TP_SIZES}, got {contract.tp_size}." + ) + if contract.device.type != "cuda": + raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.") + if torch.cuda.get_device_capability(contract.device)[0] != 10: + raise ValueError("K3 latent-MoE tail fusion requires SM100.") + if contract.dtype != torch.bfloat16: + raise ValueError("K3 latent-MoE tail fusion requires bfloat16.") + if (contract.hidden_size, contract.latent_size) != (7168, 3584): + raise ValueError( + "K3 latent-MoE tail fusion requires hidden_size=7168 and " + "latent_size=3584." + ) + + self.contract = contract + self.rank = dist.get_rank(group) + from .cute_dsl.latent_moe_tail import ( + AdaptiveUpProjectionKernel, + CollectiveKernel, + LamportCopyKernel, + ) + + with torch.accelerator.device_index(contract.device.index): + self._collective = CollectiveKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + max_token_ctas=_COLLECTIVE_TOKEN_CTAS, + rms_eps=contract.rms_eps, + fp32_internal=False, + ) + self._up_projection = AdaptiveUpProjectionKernel( + group=group, + rank=self.rank, + tp_size=contract.tp_size, + latent_dim=contract.latent_size, + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + skinny_max_m=_SKINNY_MAX_NUM_TOKENS, + mma_tiler_mn=_MMA_TILER_MN, + cluster_shape_mn=_GEMM_CLUSTER_MN, + b_prime_stages=_B_PRIME_STAGES, + ) + self._lamport_copy = LamportCopyKernel( + hidden_dim=contract.hidden_size, + max_m=contract.max_num_tokens, + ctas=_LAMPORT_COPY_CTAS, + threads=_LAMPORT_COPY_THREADS, + ) + register_cutedsl_warmup_provider(self) + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + contract = self.contract + skinny_units = tuple( + CuTeDSLCompileUnit( + name="K3 latent MoE tail Skinny up-projection", + key=( + "k3-latent-moe-tail-skinny-up-projection", + contract, + m, + ), + compile=partial(self._up_projection.compile_skinny, m), + ) + for m in range(1, self._up_projection.skinny_max_m + 1) + ) + return skinny_units + ( + CuTeDSLCompileUnit( + name="K3 latent MoE tail dynamic up-projection", + key=( + "k3-latent-moe-tail-dynamic-up-projection", + contract, + _MMA_TILER_MN, + _GEMM_CLUSTER_MN, + _B_PRIME_STAGES, + ), + compile=self._up_projection.compile_dynamic, + ), + ) + + def __call__( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> torch.Tensor: + self._validate_inputs( + routed_output, + shared_output, + rms_weight, + up_weight, + ) + self._up_projection.ensure_compiled(routed_output.shape[0]) + latent, shared_shard = self._collective( + routed_output, + shared_output, + rms_weight, + ) + local_hidden_size = self.contract.hidden_size // self.contract.tp_size + local_up_weight = up_weight.narrow( + 0, + self.rank * local_hidden_size, + local_hidden_size, + ) + mailbox = self._up_projection( + latent, + local_up_weight, + shared_shard, + ) + return self._lamport_copy( + mailbox, + m=routed_output.shape[0], + ).squeeze(0) + + def _validate_inputs( + self, + routed_output: torch.Tensor, + shared_output: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + ) -> None: + contract = self.contract + if routed_output.ndim != 2: + raise ValueError("routed_output must be a 2D tensor.") + num_tokens = routed_output.shape[0] + if routed_output.shape != (num_tokens, contract.latent_size): + raise ValueError( + f"routed_output must have shape [M, {contract.latent_size}]." + ) + if shared_output.shape != (num_tokens, contract.hidden_size): + raise ValueError( + f"shared_output must have shape [M, {contract.hidden_size}]." + ) + if rms_weight.shape != (contract.latent_size,): + raise ValueError(f"rms_weight must have shape [{contract.latent_size}].") + if up_weight.shape != (contract.hidden_size, contract.latent_size): + raise ValueError( + "up_weight must have shape " + f"[{contract.hidden_size}, {contract.latent_size}]." + ) + if not 1 <= num_tokens <= contract.max_num_tokens: + raise ValueError( + "K3 latent-MoE tail fusion requires between 1 and " + f"{contract.max_num_tokens} tokens." + ) + + tensors = (routed_output, shared_output, rms_weight, up_weight) + if any(tensor.device != contract.device for tensor in tensors): + raise ValueError("All inputs must be on the contract device.") + if any(tensor.dtype != contract.dtype for tensor in tensors): + raise ValueError("All inputs must use the contract dtype.") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("All inputs must be contiguous.") diff --git a/vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py b/vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py new file mode 100644 index 00000000000..bb577e1ba19 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/sequence_parallel.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) + + +def _custom_collective( + name: str, + x: torch.Tensor, +) -> torch.Tensor | None: + device_communicator = get_tp_group().device_communicator + if device_communicator is None: + return None + collective = getattr(device_communicator, name, None) + if collective is None: + return None + return collective(x) + + +def sp_all_gather(x: torch.Tensor) -> torch.Tensor: + output = _custom_collective("custom_all_gather", x) + if output is not None: + return output + return tensor_model_parallel_all_gather(x, 0) + + +def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + output = _custom_collective("custom_reduce_scatter", x) + if output is not None: + return output + return tensor_model_parallel_reduce_scatter(x, 0) + + +def sp_shard(x: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + sp_pad = (-x.shape[0]) % tp_size + if sp_pad > 0: + x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad)) + chunk = x.shape[0] // tp_size + return x[tp_rank * chunk : (tp_rank + 1) * chunk] + + +def sp_padding_mask( + is_padding: torch.Tensor | None, + hidden_states: torch.Tensor, +) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + if is_padding is None: + is_padding = hidden_states.new_zeros(num_tokens, dtype=torch.bool) + assert is_padding.shape[0] == num_tokens + + tp_size = get_tensor_model_parallel_world_size() + sp_pad = (-num_tokens) % tp_size + if sp_pad > 0: + is_padding = torch.nn.functional.pad(is_padding, (0, sp_pad), value=True) + chunk = is_padding.shape[0] // tp_size + tp_rank = get_tensor_model_parallel_rank() + return is_padding[tp_rank * chunk : (tp_rank + 1) * chunk] diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py b/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py new file mode 100644 index 00000000000..92936016ab6 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .chunk import ( + chunk_kda, + chunk_kda_fwd, + chunk_kda_with_fused_gate, + chunk_kda_with_fused_gate_fwd, + fused_kda_gate, + fused_kda_gate_chunk_cumsum, +) +from .fused_recurrent import ( + fused_recurrent_kda, + fused_recurrent_kda_fwd, + fused_recurrent_kda_packed_decode, +) + +__all__ = [ + "chunk_kda", + "chunk_kda_fwd", + "chunk_kda_with_fused_gate", + "chunk_kda_with_fused_gate_fwd", + "fused_kda_gate", + "fused_kda_gate_chunk_cumsum", + "fused_recurrent_kda", + "fused_recurrent_kda_fwd", + "fused_recurrent_kda_packed_decode", +] diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py new file mode 100644 index 00000000000..01bc6091289 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -0,0 +1,938 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + + +import torch + +from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h, +) +from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.op import exp2, log +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 + +from .chunk_intra import chunk_kda_fwd_intra + +BT_LIST_AUTOTUNE = [32, 64, 128] +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] + + +@triton.heuristics( + { + "STORE_QG": lambda args: args["qg"] is not None, + "STORE_KG": lambda args: args["kg"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + q, + k, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_A = tl.load(p_A, boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_k = tl.make_block_ptr( + k + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_b[:, None] + + p_gk = tl.make_block_ptr( + gk + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb *= exp2(b_gk) + if STORE_QG: + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_qg = tl.make_block_ptr( + qg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + if STORE_KG: + last_idx = min(i_t * BT + BT, T) - 1 + + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load( + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0 + ) + b_kg = b_k * exp2(b_gn - b_gk) + + p_kg = tl.make_block_ptr( + kg + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + q: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) if gk is not None else None + recompute_w_u_fwd_kernel[(NT, B * H)]( + q=q, + k=k, + qg=None, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + DOT_PRECISION="ieee", + ) + return w, u, None, kg + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gla_fwd_kernel_o( + q, + v, + g, + h, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_g = tl.make_block_ptr( + g + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_h = tl.make_block_ptr( + h + (i_tg * H + i_h) * K * V, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), + (1, 0), + ) + + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = (b_q * scale).to(b_q.dtype) + # [BT, BK] + b_g = tl.load(p_g, boundary_check=(0, 1)) + # [BT, BK] + b_qg = (b_q * exp2(b_g)).to(b_q.dtype) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + # [BT, BV] + if i_k >= 0: + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + # [BT, BV] + b_v = tl.load(p_v, boundary_check=(0, 1)) + # [BT, BT] + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) + b_o += tl.dot(b_A, b_v, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gla_fwd_o_gk( + q: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + o: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, +): + B, T, H, K, V = *q.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + def grid(meta): + return (cdiv(V, meta["BV"]), NT, B * H) + + chunk_gla_fwd_kernel_o[grid]( + q=q, + v=v, + g=g, + h=h, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + return o + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["g_bias"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in [32, 64] + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + raw_beta, + A_log, + g_bias, + o, + beta_out, + cu_seqlens, + chunk_indices, + cumsum_scale, + lower_bound, + beta, + threshold, + T, + stride_beta_batch, + stride_beta_token, + stride_beta_head, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos = i_b * T + + if i_s == 0: + o_beta_t = tl.arange(0, BT) + m_beta = i_t * BT + o_beta_t < T + if IS_VARLEN: + p_beta = ( + raw_beta + + (bos + i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + else: + p_beta = ( + raw_beta + + i_b * stride_beta_batch + + (i_t * BT + o_beta_t) * stride_beta_token + + i_h * stride_beta_head + ) + b_beta = tl.load(p_beta, mask=m_beta, other=0.0).to(tl.float32) + p_beta_out = beta_out + (bos + i_t * BT + o_beta_t) * H + i_h + tl.store(p_beta_out, tl.sigmoid(b_beta), mask=m_beta) + return + + i_s -= 1 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + if HAS_BIAS: + p_bias = tl.make_block_ptr( + g_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32) + b_s += b_bias[None, :] + + b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_s) + else: + b_g_scaled = b_s * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_s, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus + + # Boundary loads return zero, but bias and gate activation can make padded + # rows nonzero. Padding trails valid rows, so it only affects masked stores. + b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate_chunk_cumsum( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + output_dtype: torch.dtype | None = torch.float, +) -> tuple[torch.Tensor, torch.Tensor]: + if cu_seqlens is not None: + assert raw_g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + B, T, H, D = raw_g.shape + if raw_beta.shape != (B, T, H): + raise ValueError( + f"Expected raw_beta shape {(B, T, H)}, got {raw_beta.shape}" + ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) + + A_log = A_log.reshape(-1) + if g_bias is not None: + g_bias = g_bias.reshape(-1) + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype) + beta_out = torch.empty(raw_beta.shape, device=raw_beta.device, dtype=torch.float32) + + def grid(meta): + # For each (chunk, head), program 0 computes beta without extending a + # gate tile's critical path. The remaining programs cover the gate dim. + return (cdiv(meta["S"], meta["BS"]) + 1, NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + o=y, + beta_out=beta_out, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream + # exp2-based kernels reproduce exp(g). Keep this in sync with the + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. + cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, + beta=beta, + threshold=threshold, + T=T, + stride_beta_batch=raw_beta.stride(0), + stride_beta_token=raw_beta.stride(1), + stride_beta_head=raw_beta.stride(2), + H=H, + S=D, + BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, + ) + return y, beta_out + + +def _chunk_kda_fwd_with_cumulative_g( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, + safe_gate: bool = False, +): + # `g` must already be chunk-local cumulatively-summed AND scaled by + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of + # calling this helper directly unless that invariant is upheld. + Aqk, A = chunk_kda_fwd_intra( + q=q, + k=k, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + w, u, _, kg = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + del A + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, kg + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + o=v, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + del Aqk, v_new, h + return o, final_state + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log + # space so exp(x) is preserved as exp2(x / ln(2)). + g = g * RCP_LN2 + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_with_fused_gate_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + lower_bound: float | None = None, + cu_seqlens: torch.Tensor | None = None, +): + chunk_size = FLA_CHUNK_SIZE + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + g, beta = fused_kda_gate_chunk_cumsum( + raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + return _chunk_kda_fwd_with_cumulative_g( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=lower_bound is not None, + ) + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_fwd( + q=q, + k=k, + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state.contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +def chunk_kda_with_fused_gate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + g_bias: torch.Tensor | None, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + lower_bound: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + **kwargs, +): + """Run chunk KDA from raw gate and beta projections.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + + o, final_state = chunk_kda_with_fused_gate_fwd( + q=q, + k=k, + v=v.contiguous(), + raw_g=raw_g.contiguous(), + raw_beta=raw_beta, + A_log=A_log, + g_bias=g_bias, + scale=scale, + initial_state=initial_state.contiguous() if initial_state is not None else None, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + +@triton.autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) + for bt in BT_LIST_AUTOTUNE + for nw in NUM_WARPS_AUTOTUNE + for ns in [2, 3] + ], + key=["H", "D"], +) +@triton.jit +def kda_gate_fwd_kernel( + g, + A, + y, + g_bias, + lower_bound, + beta: tl.constexpr, + threshold: tl.constexpr, + T, + H, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_t, i_h = tl.program_id(0), tl.program_id(1) + n_t = i_t * BT + + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + + stride_row = H * D + stride_col = 1 + + g_ptr = tl.make_block_ptr( + base=g + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + y_ptr = tl.make_block_ptr( + base=y + i_h * D, + shape=(T, D), + strides=(stride_row, stride_col), + offsets=(n_t, 0), + block_shape=(BT, BD), + order=(1, 0), + ) + + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + n_d = tl.arange(0, BD) + bias_mask = n_d < D + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( + tl.float32 + ) + b_g = b_g + b_bias[None, :] + + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp + + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) + + +def fused_kda_gate( + g: torch.Tensor, + A: torch.Tensor, + head_k_dim: int, + g_bias: torch.Tensor | None = None, + beta: float = 1.0, + threshold: float = 20.0, + lower_bound: float | None = None, +) -> torch.Tensor: + """ + Forward pass for KDA gate: + input g: [..., H*D] + param A: [H] or [1, 1, H, 1] + beta: softplus beta parameter + threshold: softplus threshold parameter + return : [..., H, D] + """ + orig_shape = g.shape[:-1] + + g = g.view(-1, g.shape[-1]) + T = g.shape[0] + HD = g.shape[1] + H = A.numel() + assert H * head_k_dim == HD + assert g.stride() == (HD, 1) + + y = torch.empty_like(g, dtype=torch.float32) + + def grid(meta): + return (cdiv(T, meta["BT"]), H) + + kda_gate_fwd_kernel[grid]( + g, + A, + y, + g_bias, + lower_bound or 0.0, + beta, + threshold, + T, + H, + head_k_dim, + BD=next_power_of_2(head_k_dim), + HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, + ) + + y = y.view(*orig_shape, H, head_k_dim) + return y diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py new file mode 100644 index 00000000000..087197d10a5 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra.py @@ -0,0 +1,559 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices +from vllm.third_party.flash_linear_attention.ops.op import exp2, gather +from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported +from vllm.triton_utils import tl, triton + +from .chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=["H", "HV", "K", "BC"], +) +@triton.jit(do_not_specialize=['T']) +def chunk_kda_fwd_kernel_inter_solve_fused( + q, + k, + g, + beta, + Aqk, + Akkd, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_SAFE_GATE: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akkd (fp32) + 4. Does forward substitution on diagonals + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + g += (bos * HV + i_hv) * K + Aqk += (bos * HV + i_hv) * BT + Akk += (bos * HV + i_hv) * BT + Akkd += (bos * HV + i_hv) * BC + + o_i = tl.arange(0, BC) + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # off-diagonal blocks + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) + b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) + b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + + if i_tc1 < T: + p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) + b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) + b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + # [BC, BC] + b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) + b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + + if i_tc2 < T: + p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) + b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) + b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_Aqk20 += tl.dot(b_qg2, b_kgt) + b_Akk20 += tl.dot(b_kg2, b_kgt) + # [BC, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + # [BC, BC] + b_Aqk21 += tl.dot(b_qg2, b_kgt) + b_Akk21 += tl.dot(b_kg2, b_kgt) + + if i_tc3 < T: + p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + # [BC, BK] + b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) + b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) + b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + # [BK] + b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + # [BK, BC] + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + # [BC, BC] + b_Aqk30 += tl.dot(b_qg3, b_kgt) + b_Akk30 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + # [BC, BC] + b_Aqk31 += tl.dot(b_qg3, b_kgt) + b_Akk31 += tl.dot(b_kg3, b_kgt) + # [BK, BC] + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + # [BC, BC] + b_Aqk32 += tl.dot(b_qg3, b_kgt) + b_Akk32 += tl.dot(b_kg3, b_kgt) + + ################################################################################ + # save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) + b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) + b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + + p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) + b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + + ################################################################################ + # forward substitution on diagonals + ################################################################################ + + if not USE_SAFE_GATE: + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # store full Akk_inv to Akk + ################################################################################ + + p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) + p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) + p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) + p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) + p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) + p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) + p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) + p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) + p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT', 'HV'], +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_kda_fwd_kernel_intra_sub_chunk( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + g = g + (bos * HV + i_hv) * K + beta = beta + bos * HV + i_hv + Aqk = Aqk + (bos * HV + i_hv) * BT + Akk = Akk + (bos * HV + i_hv) * BC + + p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) + + p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + + b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # caculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + b_kgt = tl.trans(b_k * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) + p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, +): + B, T, H, K, HV = *k.shape, gk.shape[2] + BT = chunk_size + BC = 16 + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) + + # Compute diagonal blocks into Akkd in fp32. + if safe_gate: + grid = (NT, NC, B * HV) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) + solve_tril_dot_precision = ( + "tf32" + if current_platform.is_cuda() + and current_platform.has_device_capability(80) + else "ieee" + ) + grid = (NT, B * HV) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akkd=Akkd, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + USE_SAFE_GATE=safe_gate, + SOLVE_TRIL_DOT_PRECISION=solve_tril_dot_precision, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py new file mode 100644 index 00000000000..ecd00a51f9d --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk_intra_token_parallel.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# Forward-only adaptation of flash-linear-attention 0.5.0. +# ruff: noqa: E501 + +# Token-parallel implementation of KDA intra chunk kernel + +import torch + +from vllm.third_party.flash_linear_attention.ops.op import exp2 +from vllm.triton_utils import tl, triton + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H", "HV"], +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_kda_fwd_kernel_intra_token_parallel( + q, + k, + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_tg, i_hg = tl.program_id(0), tl.program_id(1) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + G: tl.constexpr = HV // H + + q += bos * H*K + k += bos * H*K + g += bos * HV*K + Aqk += bos * HV*BT + Akk += bos * HV*BC + beta += bos * HV + + BK: tl.constexpr = triton.next_power_of_2(K) + o_hv = i_hg * BH + tl.arange(0, BH) + o_h = o_hv // G + o_k = tl.arange(0, BK) + m_hv = o_hv < HV + m_k = o_k < K + m_hk = m_hv[:, None] & m_k[None, :] + + # q/k: [B, T, H, K], manual load via mapped qk head index + p_qk = o_h[:, None] * K + o_k[None, :] + b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + + # g: [B, T, HV, K], beta: [B, T, HV] + p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) + b_k *= b_beta[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) + p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + + b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) + b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale + b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv) + tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) + + +def chunk_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation: each token gets its own thread block. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) + beta: [B, T, HV] + Aqk: [B, T, HV, BT] output tensor to write to + Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + chunk_size: BT (default 64) + sub_chunk_size: BC (default 16) + """ + B, T, H, K, HV = *q.shape, gk.shape[2] + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) + chunk_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py new file mode 100644 index 00000000000..b4d35d85571 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This file contains code adapted from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# ruff: noqa: E501 + +import torch + +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops.op import exp, log +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + + +@triton.heuristics( + { + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit +def _kda_gate_beta_fwd_kernel( + raw_g, + raw_beta, + A_log, + dt_bias, + gate, + beta_out, + lower_bound, + softplus_beta: tl.constexpr, + softplus_threshold: tl.constexpr, + T, + stride_g_token: tl.constexpr, + stride_beta_token: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + i_t, i_h = tl.program_id(0), tl.program_id(1) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_d = o_d < D + + p_g = raw_g + o_t[:, None] * stride_g_token + i_h * D + o_d[None, :] + b_g = tl.load(p_g, mask=m_t[:, None] & m_d[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * D + o_d, + mask=m_d, + other=0.0, + ).to(tl.float32) + b_g += b_bias[None, :] + + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_scaled = b_g * softplus_beta + b_softplus = tl.where( + b_scaled > softplus_threshold, + b_g, + log(1.0 + tl.exp(b_scaled)) / softplus_beta, + ) + b_gate = -b_a * b_softplus + + p_gate = gate + (o_t[:, None] * H + i_h) * D + o_d[None, :] + tl.store( + p_gate, + b_gate, + mask=m_t[:, None] & m_d[None, :], + ) + + b_beta = tl.load( + raw_beta + o_t * stride_beta_token + i_h, + mask=m_t, + other=0.0, + ).to(tl.float32) + tl.store(beta_out + o_t * H + i_h, tl.sigmoid(b_beta), mask=m_t) + + +def _fused_kda_gate_beta( + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, D = raw_g.shape + assert B == 1 + assert raw_beta.shape == (B, T, H) + assert raw_g.stride()[2:] == (D, 1) + assert raw_beta.stride(2) == 1 + gate = torch.empty((B, T, H, D), dtype=torch.float32, device=raw_g.device) + beta = torch.empty((B, T, H), dtype=torch.float32, device=raw_beta.device) + + BT = 16 + _kda_gate_beta_fwd_kernel[(cdiv(T, BT), H)]( + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + gate=gate, + beta_out=beta, + lower_bound=lower_bound, + softplus_beta=1.0, + softplus_threshold=20.0, + T=T, + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + H=H, + D=D, + BT=BT, + BD=next_power_of_2(D), + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return gate, beta + + +@triton.heuristics( + { + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T", "stride_beta_token"]) +def fused_recurrent_kda_fwd_kernel( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + out, + state, + cu_seqlens, + state_indices, + num_accepted_tokens, + lower_bound, + scale: tl.constexpr, + N: tl.int64, + T: tl.int64, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_qkv_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_out_token: tl.constexpr, + stride_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + num_stages: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + i_v = pid % tl.cdiv(V, BV) + i_nh = pid // tl.cdiv(V, BV) + i_n, i_h = i_nh // H, i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + sequence_length = eos - bos + if sequence_length == 0: + return + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_state = m_v[:, None] & m_k[None, :] + + if IS_SPEC_DECODING: + initial_token = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + initial_token = 0 + state_index = tl.load(state_indices + i_n * stride_indices_seq + initial_token).to( + tl.int64 + ) + p_out = out + bos * stride_out_token + i_h * V + o_v + if state_index <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=m_v) + return + + p_state = ( + state + + state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + b_state = tl.load(p_state, mask=m_state, other=0.0).to(tl.float32) + + p_q = q + bos * stride_qkv_token + i_h * K + o_k + p_k = k + bos * stride_qkv_token + i_h * K + o_k + p_v = v + bos * stride_qkv_token + i_h * V + o_v + p_g = g + bos * stride_g_token + i_h * K + o_k + p_beta = beta + bos * stride_beta_token + i_h + for i_t in tl.range(0, sequence_length, num_stages=num_stages): + b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_k = tl.load(p_k, mask=m_k, other=0.0, eviction_policy="evict_last").to( + tl.float32 + ) + b_v = tl.load(p_v, mask=m_v, other=0.0, eviction_policy="evict_first").to( + tl.float32 + ) + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + b_gate = tl.load( + p_g, + mask=m_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + if USE_GATE_IN_KERNEL: + if HAS_DT_BIAS: + b_bias = tl.load( + dt_bias + i_h * K + o_k, + mask=m_k, + other=0.0, + ).to(tl.float32) + b_gate += b_bias + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_gate) + else: + b_softplus = tl.where( + b_gate > 20.0, + b_gate, + log(1.0 + tl.exp(b_gate)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + tl.store( + p_out, + b_out.to(p_out.dtype.element_ty), + mask=m_v, + eviction_policy="evict_first", + ) + + final_state_index = tl.load(state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + if final_state_index > 0: + p_final_state = ( + state + + final_state_index * stride_state_token + + i_h * V * K + + o_v[:, None] * K + + o_k[None, :] + ) + tl.store( + p_final_state, + b_state.to(p_final_state.dtype.element_ty), + mask=m_state, + ) + + p_q += stride_qkv_token + p_k += stride_qkv_token + p_v += stride_qkv_token + p_g += stride_g_token + p_beta += stride_beta_token + p_out += stride_out_token + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_fused_recurrent_kda_fwd_warmup_profiles( + num_heads: int, +) -> tuple[int, ...]: + """Return representative sequence counts for gated launch variants.""" + # The region above 192 head-sequences reuses the second launch variant. + return ( + 1, + 48 // num_heads + 1, + 96 // num_heads + 1, + ) + + +def fused_recurrent_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + inplace_final_state: bool = True, + cu_seqlens: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + use_gate_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch recurrent KDA with dense inner dimensions and row strides.""" + B, T, H, K = q.shape + V = v.shape[-1] + assert B == 1 and k.shape == q.shape + assert v.shape == (B, T, H, V) and g.shape == (B, T, H, K) + assert beta.shape == (B, T, H) + assert initial_state is not None + assert cu_seqlens is not None + assert ssm_state_indices is not None + assert inplace_final_state + if out is None: + out = torch.empty_like(v) + assert out.shape == v.shape + assert initial_state.shape[1:] == (H, V, K) + assert ssm_state_indices.ndim in (1, 2) + + assert q.stride()[2:] == k.stride()[2:] == (K, 1) + assert v.stride()[2:] == out.stride()[2:] == (V, 1) + assert g.stride()[2:] == (K, 1) + assert beta.stride(2) == 1 + assert q.stride(1) == k.stride(1) == v.stride(1) + assert initial_state.stride()[1:] == (V * K, K, 1) + N = cu_seqlens.numel() - 1 + if ssm_state_indices.ndim == 1: + assert T == N + assert num_accepted_tokens is None + else: + assert ssm_state_indices.stride(1) == 1 + assert cu_seqlens.is_contiguous() + if use_gate_in_kernel: + assert A_log is not None and A_log.is_contiguous() + assert dt_bias is None or dt_bias.is_contiguous() + + if scale is None: + scale = K**-0.5 + + if use_gate_in_kernel: + # Tuned on GB300 for Kimi-K3 shapes. Keep the warmup profiles above in + # sync with these boundaries. + head_sequences = H * N + if head_sequences <= 48: + BV, num_stages = 4, 4 + elif head_sequences <= 96: + BV, num_stages = 8, 3 + elif head_sequences <= 192: + BV, num_stages = 16, 3 + else: + BV, num_stages = 8, 3 + num_warps = 1 + else: + BV, num_warps, num_stages = 8, 1, 2 + grid = (cdiv(V, BV) * N * H,) + fused_recurrent_kda_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + cu_seqlens=cu_seqlens, + state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + lower_bound=lower_bound, + scale=scale, + N=N, + T=T, + H=H, + K=K, + V=V, + BK=next_power_of_2(K), + BV=BV, + stride_qkv_token=q.stride(1), + stride_g_token=g.stride(1), + stride_beta_token=beta.stride(1), + stride_out_token=out.stride(1), + stride_state_token=initial_state.stride(0), + stride_indices_seq=ssm_state_indices.stride(0), + IS_SPEC_DECODING=num_accepted_tokens is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + USE_GATE_IN_KERNEL=use_gate_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + num_warps=num_warps, + num_stages=num_stages, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state + + +def fused_recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor | None = None, + out: torch.Tensor | None = None, + fuse_gate: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run recurrent KDA from raw gate and beta inputs. + + This vLLM wrapper applies the gate activation and beta sigmoid, selecting + whether to materialize them before launching the recurrent kernel. + """ + if fuse_gate is None: + fuse_gate = True + + if fuse_gate: + gate = raw_g + beta = raw_beta + else: + gate, beta = _fused_kda_gate_beta( + raw_g, + raw_beta, + A_log, + dt_bias, + lower_bound, + ) + return fused_recurrent_kda_fwd( + q=q, + k=k, + v=v, + g=gate, + beta=beta, + scale=q.shape[-1] ** -0.5, + initial_state=initial_state, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log if fuse_gate else None, + dt_bias=dt_bias if fuse_gate else None, + lower_bound=lower_bound if fuse_gate else None, + use_gate_in_kernel=fuse_gate, + use_beta_sigmoid_in_kernel=fuse_gate, + out=out, + ) + + +@triton.jit( + do_not_specialize=["stride_beta_token", "stride_state_indices"] +) +def fused_recurrent_kda_packed_decode_kernel( + mixed_qkv, + raw_g, + raw_beta, + A_log, + dt_bias, + out, + state, + state_indices, + lower_bound, + scale: tl.constexpr, + stride_mixed_token: tl.constexpr, + stride_g_token: tl.constexpr, + stride_beta_token, + stride_state_token: tl.constexpr, + stride_state_indices, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + state_idx = tl.load(state_indices + i_n * stride_state_indices).to(tl.int64) + p_out = out + (i_n * H + i_h) * V + o_v + if state_idx <= 0: + tl.store(p_out, tl.zeros([BV], dtype=tl.float32), mask=mask_v) + return + + p_state = state + state_idx * stride_state_token + p_state += i_h * V * K + o_v[:, None] * K + o_k[None, :] + b_state = tl.load(p_state, mask=mask_state, other=0).to(tl.float32) + + # Q, K, and V occupy consecutive channel ranges, while the token stride + # may also include the output-gate projection that follows packed QKV. + p_mixed = mixed_qkv + i_n * stride_mixed_token + b_q = tl.load(p_mixed + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load( + p_mixed + H * K + i_h * K + o_k, + mask=mask_k, + other=0, + ).to(tl.float32) + b_v = tl.load( + p_mixed + 2 * H * K + i_h * V + o_v, + mask=mask_v, + other=0, + ).to(tl.float32) + + b_q /= tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k /= tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q *= scale + + p_g = raw_g + i_n * stride_g_token + i_h * K + o_k + b_g = tl.load(p_g, mask=mask_k, other=0).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * K + o_k, mask=mask_k, other=0).to(tl.float32) + b_a = exp(tl.load(A_log + i_h).to(tl.float32)) + b_g += b_bias + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_softplus = tl.where( + b_g > SOFTPLUS_THRESHOLD, + b_g, + log(1.0 + tl.exp(b_g)), + ) + b_gate = -b_a * b_softplus + + b_state *= exp(b_gate[None, :]) + b_v -= tl.sum(b_state * b_k[None, :], axis=1) + b_beta = tl.sigmoid( + tl.load(raw_beta + i_n * stride_beta_token + i_h).to(tl.float32) + ) + b_v *= b_beta + b_state += b_v[:, None] * b_k[None, :] + b_out = tl.sum(b_state * b_q[None, :], axis=1) + + tl.store(p_out, b_out.to(p_out.dtype.element_ty), mask=mask_v) + tl.store(p_state, b_state.to(p_state.dtype.element_ty), mask=mask_state) + + +def fused_recurrent_kda_packed_decode( + mixed_qkv: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + initial_state: torch.Tensor, + state_indices: torch.Tensor, + scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run one-token KDA decode directly from packed post-conv QKV.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("`mixed_qkv` must be 2D and contiguous in its last dim.") + if raw_g.ndim != 4 or raw_g.shape[0] != 1: + raise ValueError("`raw_g` must have shape [1, B, H, K].") + if raw_beta.ndim != 3 or raw_beta.shape[0] != 1: + raise ValueError("`raw_beta` must have shape [1, B, H].") + if initial_state.ndim != 4: + raise ValueError("`initial_state` must have shape [cache, H, V, K].") + _, H, V, K = initial_state.shape + if raw_g.stride()[2:] != (K, 1): + raise ValueError("`raw_g` must be contiguous within each token.") + if raw_beta.stride(2) != 1: + raise ValueError("`raw_beta` heads must be contiguous.") + if initial_state.stride()[1:] != (V * K, K, 1): + raise ValueError("`initial_state` must be contiguous within each cache slot.") + if state_indices.ndim != 1: + raise ValueError("`state_indices` must be one-dimensional.") + if A_log.ndim != 1 or not A_log.is_contiguous(): + raise ValueError("`A_log` must be contiguous and one-dimensional.") + if not dt_bias.is_contiguous(): + raise ValueError("`dt_bias` must be contiguous.") + + device = mixed_qkv.device + if any( + x.device != device + for x in (raw_g, raw_beta, A_log, dt_bias, initial_state, state_indices) + ): + raise ValueError("All packed KDA inputs must be on the same device.") + + B = mixed_qkv.shape[0] + if raw_g.shape != (1, B, H, K): + raise ValueError(f"Unexpected raw gate shape {tuple(raw_g.shape)}.") + if raw_beta.shape != (1, B, H): + raise ValueError(f"Unexpected raw beta shape {tuple(raw_beta.shape)}.") + if mixed_qkv.shape[1] != 2 * H * K + H * V: + raise ValueError(f"Unexpected packed QKV shape {tuple(mixed_qkv.shape)}.") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise ValueError("`A_log` or `dt_bias` has an incompatible shape.") + if state_indices.shape[0] != B: + raise ValueError("`state_indices` must contain one entry per token.") + + BK = next_power_of_2(K) + BV = min(next_power_of_2(V), 32) + if scale is None: + scale = K**-0.5 + + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + grid = (cdiv(V, BV), B * H) + fused_recurrent_kda_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + out=out, + state=initial_state, + state_indices=state_indices, + lower_bound=lower_bound or 0.0, + scale=scale, + stride_mixed_token=mixed_qkv.stride(0), + stride_g_token=raw_g.stride(1), + stride_beta_token=raw_beta.stride(1), + stride_state_token=initial_state.stride(0), + stride_state_indices=state_indices.stride(0), + H=H, + K=K, + V=V, + BK=BK, + BV=BV, + SOFTPLUS_THRESHOLD=20.0, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, initial_state diff --git a/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py b/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py new file mode 100644 index 00000000000..3d429a9b58e --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/vision_fa4_warmup.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup compilation of Kimi-K3 vision FA4 kernels.""" + +from __future__ import annotations + +import math +from collections.abc import Iterator +from dataclasses import dataclass +from functools import partial + +import torch + +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) +from vllm.platforms import current_platform + +_FA4_TILE_SIZE = 128 +_FA4_MAX_SPLITS = 128 + + +@dataclass(frozen=True) +class KimiK3VisionFA4WarmupConfig: + num_heads: int + head_dim: int + dtype: torch.dtype + max_batch_size: int + max_seqlen: int + + +@dataclass(frozen=True) +class _FA4WarmupProbe: + batch_size: int + max_seqlen: int + dispatch_key: tuple[int, bool, int | None] + + +def _combine_log_max_splits(head_dim: int, num_splits: int) -> int: + k_block_size = 64 if head_dim <= 64 else 128 + tile_m = 8 if k_block_size == 128 else 16 + min_log_max_splits = 5 if tile_m == 8 else 4 + return max(math.ceil(math.log2(num_splits)), min_log_max_splits) + + +def _get_dispatch( + config: KimiK3VisionFA4WarmupConfig, + *, + batch_size: int, + max_seqlen: int, + num_sms: int, +) -> tuple[int, bool, int | None]: + # These are the shape-dependent fields in FA4's forward and combine + # compile keys for noncausal varlen MHA. Compile units still pass + # num_splits=0 so the production selector makes the final decision. + q_stage = 2 if max_seqlen > _FA4_TILE_SIZE else 1 + effective_q_tile = q_stage * _FA4_TILE_SIZE + num_m_blocks = math.ceil(max_seqlen / effective_q_tile) + num_n_blocks = math.ceil(max_seqlen / _FA4_TILE_SIZE) + + if num_n_blocks <= 4: + num_splits = 1 + else: + total_m_blocks = batch_size * config.num_heads * num_m_blocks + num_splits = min( + num_sms // total_m_blocks, + _FA4_MAX_SPLITS, + num_n_blocks, + ) + + is_split_kv = num_splits > 1 + combine_key = ( + _combine_log_max_splits(config.head_dim, num_splits) if is_split_kv else None + ) + return q_stage, is_split_kv, combine_key + + +def _get_warmup_probes( + config: KimiK3VisionFA4WarmupConfig, + *, + num_sms: int, +) -> tuple[_FA4WarmupProbe, ...]: + if config.max_batch_size <= 0 or config.max_seqlen <= 0: + return () + + probes: dict[tuple[int, bool, int | None], _FA4WarmupProbe] = {} + + def add_probe(batch_size: int, max_seqlen: int) -> None: + dispatch_key = _get_dispatch( + config, + batch_size=batch_size, + max_seqlen=max_seqlen, + num_sms=num_sms, + ) + if dispatch_key not in probes: + probes[dispatch_key] = _FA4WarmupProbe( + batch_size=batch_size, + max_seqlen=max_seqlen, + dispatch_key=dispatch_key, + ) + + add_probe(batch_size=1, max_seqlen=1) + if config.max_seqlen > _FA4_TILE_SIZE: + add_probe(batch_size=1, max_seqlen=_FA4_TILE_SIZE + 1) + + max_n_blocks = math.ceil(config.max_seqlen / _FA4_TILE_SIZE) + for num_n_blocks in range(5, max_n_blocks + 1): + max_seqlen = min( + (num_n_blocks - 1) * _FA4_TILE_SIZE + 1, + config.max_seqlen, + ) + num_m_blocks = math.ceil(max_seqlen / (2 * _FA4_TILE_SIZE)) + max_split_batch_size = min( + config.max_batch_size, + num_sms // (2 * config.num_heads * num_m_blocks), + ) + if max_split_batch_size == 0: + break + for batch_size in range(1, max_split_batch_size + 1): + add_probe(batch_size, max_seqlen) + + return tuple(probes.values()) + + +def _iter_compile_units( + config: KimiK3VisionFA4WarmupConfig, +) -> Iterator[CuTeDSLCompileUnit]: + device_id = current_platform.current_device() + num_sms = current_platform.num_compute_units(device_id) + for probe in _get_warmup_probes(config, num_sms=num_sms): + yield CuTeDSLCompileUnit( + name="kimi_k3_vision_fa4", + key=("kimi_k3_vision_fa4", config, probe.dispatch_key), + compile=partial(_compile, config, probe), + ) + + +def _compile( + config: KimiK3VisionFA4WarmupConfig, + probe: _FA4WarmupProbe, +) -> None: + from vllm.vllm_flash_attn import flash_attn_varlen_func + + device = current_platform.current_device() + total_tokens = probe.batch_size * probe.max_seqlen + qkv = torch.empty( + 3, + total_tokens, + config.num_heads, + config.head_dim, + device=device, + dtype=config.dtype, + ) + cu_seqlens = torch.arange( + 0, + total_tokens + 1, + probe.max_seqlen, + device=qkv.device, + dtype=torch.int32, + ) + flash_attn_varlen_func( + qkv[0], + qkv[1], + qkv[2], + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=probe.max_seqlen, + max_seqlen_k=probe.max_seqlen, + dropout_p=0.0, + causal=False, + softmax_scale=config.head_dim**-0.5, + fa_version=4, + num_splits=0, + ) + + +class _WarmupProvider: + def __init__(self) -> None: + self.configs: set[KimiK3VisionFA4WarmupConfig] = set() + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + return tuple( + unit for config in self.configs for unit in _iter_compile_units(config) + ) + + +_PROVIDER = _WarmupProvider() + + +def register_kimi_k3_vision_fa4_warmup( + config: KimiK3VisionFA4WarmupConfig, +) -> None: + _PROVIDER.configs.add(config) + register_cutedsl_warmup_provider(_PROVIDER) diff --git a/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py new file mode 100644 index 00000000000..60fc5335070 --- /dev/null +++ b/vllm/third_party/flash_linear_attention/ops/fused_norm_gate.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source was licensed under the MIT license. +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch +import torch.nn as nn + +from vllm.model_executor.custom_op import CustomOp +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv, next_power_of_2 + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n, + D: tl.constexpr, # number of columns in x + BT: tl.constexpr, + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + + o_d = tl.arange(0, BD) + m_d = o_d < D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + if HAS_RESIDUAL: + p_res = tl.make_block_ptr( + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = tl.make_block_ptr( + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0) + ) + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1)) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=1) / D + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,)) + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + else: + b_xbar = tl.where(m_d[None, :], b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = ( + (b_x - b_mean[:, None]) * b_rstd[:, None] + if not IS_RMS_NORM + else b_x * b_rstd[:, None] + ) + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b[None, :] + + # swish/sigmoid output gate + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics( + { + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None, + "HAS_RESIDUAL": lambda args: args["residual"] is not None, + "HAS_WEIGHT": lambda args: args["w"] is not None, + "HAS_BIAS": lambda args: args["b"] is not None, + } +) +@triton.jit +def layer_norm_gated_fwd_kernel1( + x, # pointer to the input + g, # pointer to the gate + y, # pointer to the output + w, # pointer to the weights + b, # pointer to the biases + residual, # pointer to the residual + residual_out, # pointer to the residual + mean, # pointer to the mean + rstd, # pointer to the 1/std + eps, # epsilon to avoid division by zero + D: tl.constexpr, # number of columns in x + BD: tl.constexpr, + ACTIVATION: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + STORE_RESIDUAL_OUT: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + launch_pdl: tl.constexpr, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + g += i_t * D + if HAS_RESIDUAL: + residual += i_t * D + if STORE_RESIDUAL_OUT: + residual_out += i_t * D + + if launch_pdl: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + o_d = tl.arange(0, BD) + m_d = o_d < D + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + tl.store(residual_out + o_d, b_x, mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + tl.store(mean + i_t, b_mean) + b_xbar = tl.where(m_d, b_x - b_mean, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(m_d, b_x, 0.0) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + b_rstd = 1 / tl.sqrt(b_var + eps) + tl.store(rstd + i_t, b_rstd) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32) + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + # swish/sigmoid output gate + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + b_y = b_y * b_g * tl.sigmoid(b_g) + elif ACTIVATION == "sigmoid": + b_y = b_y * tl.sigmoid(b_g) + + # Write output + tl.store(y + o_d, b_y, mask=m_d) + + +def layer_norm_gated_fwd( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + eps: float = 1e-5, + residual: torch.Tensor = None, + out_dtype: torch.dtype = None, + residual_dtype: torch.dtype = None, + is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, +): + if residual is not None: + residual_dtype = residual.dtype + T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 + if residual is not None: + assert residual.shape == (T, D) + if weight is not None: + assert weight.shape == (D,) + if bias is not None: + assert bias.shape == (D,) + # allocate output + y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype) + if residual is not None or ( + residual_dtype is not None and residual_dtype != x.dtype + ): + residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) + else: + residual_out = None + mean = ( + torch.empty((T,), dtype=torch.float, device=x.device) + if not is_rms_norm + else None + ) + rstd = torch.empty((T,), dtype=torch.float, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + if D <= 512: + BT = 16 + layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + T=T, + H=H, + g_stride_n=g_stride_n, + D=D, + BD=BD, + BT=BT, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=8, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + else: + layer_norm_gated_fwd_kernel1[(T,)]( + x=x, + g=g, + y=y, + w=weight, + b=bias, + residual=residual, + residual_out=residual_out, + mean=mean, + rstd=rstd, + eps=eps, + D=D, + BD=BD, + ACTIVATION=activation, + IS_RMS_NORM=is_rms_norm, + num_warps=4, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + # residual_out is None if residual is None and residual_dtype == input_dtype + return y, mean, rstd, residual_out if residual_out is not None else x + + +def rms_norm_gated( + x: torch.Tensor, + g: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + activation: str = "swish", + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + eps: float = 1e-6, +): + x_shape_og = x.shape + # reshape input data into 2D tensor + x = x.contiguous().reshape(-1, x.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.contiguous().reshape(-1, residual.shape[-1]) + residual_dtype = ( + residual.dtype + if residual is not None + else (torch.float if residual_in_fp32 else None) + ) + y, _, _, residual_out = layer_norm_gated_fwd( + x=x, + g=g, + weight=weight, + bias=bias, + activation=activation, + eps=eps, + residual=residual, + residual_dtype=residual_dtype, + is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, + ) + y = y.reshape(x_shape_og) + return y if not prenorm else (y, residual_out.reshape(x_shape_og)) + + +@CustomOp.register("fused_rms_norm_gated") +class FusedRMSNormGated(CustomOp): + def __init__( + self, + hidden_size: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + activation: str = "swish", + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.activation = activation + + if self.activation not in ["swish", "silu", "sigmoid"]: + raise ValueError(f"Unsupported activation: {self.activation}") + + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward_native( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + """Decomposed PyTorch ops for torch.compile/inductor fusion.""" + # TODO(https://github.com/vllm-project/vllm/issues/36175): implement + # native residual/prenorm path and unify with RMSNormGated. + # For now, fall back to the triton kernel. + if residual is not None or prenorm: + return self.forward_cuda(x, g, residual, prenorm, residual_in_fp32) + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_normed = x_float * torch.rsqrt(variance + self.eps) + if self.weight is not None: + x_normed = x_normed * self.weight.float() + g_float = g.float() + if self.activation in ("swish", "silu"): + out = x_normed * g_float * torch.sigmoid(g_float) + else: # sigmoid + out = x_normed * torch.sigmoid(g_float) + return out.to(x.dtype) + + def forward_cuda( + self, + x: torch.Tensor, + g: torch.Tensor, + residual: torch.Tensor | None = None, + prenorm: bool = False, + residual_in_fp32: bool = False, + ) -> torch.Tensor: + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 10acc214f4d..1701cd47aca 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -166,6 +166,8 @@ def layer_norm_gated_fwd_kernel( rstd, # pointer to the 1/std eps, # epsilon to avoid division by zero T, # number of rows in x + H: tl.constexpr, # number of heads + g_stride_n: tl.constexpr, D: tl.constexpr, # number of columns in x BT: tl.constexpr, BD: tl.constexpr, @@ -221,8 +223,13 @@ def layer_norm_gated_fwd_kernel( b_y = b_y + b_b[None, :] # swish/sigmoid output gate - p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_g = (o_t // H) * g_stride_n + (o_t % H) * D + b_g = tl.load( + g + o_g[:, None] + o_d[None, :], + mask=(o_t[:, None] < T) & m_d[None, :], + other=0.0, + ).to(tl.float32) if ACTIVATION == "swish" or ACTIVATION == "silu": b_y = b_y * b_g * tl.sigmoid(b_g) elif ACTIVATION == "sigmoid": @@ -320,10 +327,15 @@ def layer_norm_gated_fwd( out_dtype: torch.dtype = None, residual_dtype: torch.dtype = None, is_rms_norm: bool = False, + H: int = 1, + g_stride_n: int | None = None, ): if residual is not None: residual_dtype = residual.dtype T, D = x.shape + if g_stride_n is None: + g_stride_n = D + assert T % H == 0 if residual is not None: assert residual.shape == (T, D) if weight is not None: @@ -349,10 +361,8 @@ def layer_norm_gated_fwd( BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) if D > BD: raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - if D <= 512: - BT = 32 + BT = 16 layer_norm_gated_fwd_kernel[(cdiv(T, BT),)]( x=x, g=g, @@ -365,12 +375,14 @@ def layer_norm_gated_fwd( rstd=rstd, eps=eps, T=T, + H=H, + g_stride_n=g_stride_n, D=D, BD=BD, BT=BT, ACTIVATION=activation, IS_RMS_NORM=is_rms_norm, - num_warps=4, + num_warps=8, ) else: layer_norm_gated_fwd_kernel1[(T,)]( @@ -408,7 +420,16 @@ def rms_norm_gated( x_shape_og = x.shape # reshape input data into 2D tensor x = x.contiguous().reshape(-1, x.shape[-1]) - g = g.contiguous().reshape(-1, g.shape[-1]) + D = x.shape[-1] + # The tiled kernel supports row-strided gates; kernel1 does not. + if D <= 512: + H = 1 if g.ndim == 2 else g.shape[-2] + g = g.view(-1, H, D) + g_stride_n = g.stride(0) + else: + g = g.contiguous() + H = 1 + g_stride_n = D if residual is not None: assert residual.shape == x_shape_og residual = residual.contiguous().reshape(-1, residual.shape[-1]) @@ -427,6 +448,8 @@ def rms_norm_gated( residual=residual, residual_dtype=residual_dtype, is_rms_norm=True, + H=H, + g_stride_n=g_stride_n, ) y = y.reshape(x_shape_og) return y if not prenorm else (y, residual_out.reshape(x_shape_og)) @@ -1187,6 +1210,7 @@ def kda_gate_cumsum_fwd_kernel( cu_seqlens, chunk_indices, cumsum_scale, + lower_bound, beta, threshold, T, @@ -1196,6 +1220,7 @@ def kda_gate_cumsum_fwd_kernel( BD: tl.constexpr, HAS_BIAS: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) i_b, i_h = i_bh // H, i_bh % H @@ -1235,14 +1260,17 @@ def kda_gate_cumsum_fwd_kernel( b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32) b_g = b_g + b_bias[None, :] - b_a = -tl.exp(tl.load(A + i_h).to(tl.float32)) - b_g_scaled = b_g * beta - b_softplus = tl.where( - b_g_scaled > threshold, - b_g, - (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), - ) - b_gate = b_a * b_softplus + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) + if USE_LOWER_BOUND: + b_gate = lower_bound * tl.sigmoid(b_a * b_g) + else: + b_g_scaled = b_g * beta + b_softplus = tl.where( + b_g_scaled > threshold, + b_g, + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)), + ) + b_gate = -b_a * b_softplus # Out-of-bounds rows (load returns 0, but softplus/bias can still make # b_gate non-zero) participate in the dot product. They only contribute to @@ -1260,6 +1288,7 @@ def fused_kda_gate_chunk_cumsum( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, @@ -1293,16 +1322,17 @@ def fused_kda_gate_chunk_cumsum( # exp2-based kernels reproduce exp(g). Keep this in sync with the # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`. cumsum_scale=RCP_LN2, + lower_bound=lower_bound or 0.0, beta=beta, threshold=threshold, T=T, H=H, D=D, BT=chunk_size, + USE_LOWER_BOUND=lower_bound is not None, ) return y - def _chunk_kda_fwd_with_cumulative_g( q: torch.Tensor, k: torch.Tensor, @@ -1424,6 +1454,7 @@ def chunk_kda_with_fused_gate_fwd( scale: float, initial_state: torch.Tensor, output_final_state: bool, + lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE @@ -1439,6 +1470,7 @@ def chunk_kda_with_fused_gate_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, + lower_bound=lower_bound, ) return _chunk_kda_fwd_with_cumulative_g( q=q, @@ -1500,6 +1532,7 @@ def chunk_kda_with_fused_gate( scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, + lower_bound: float | None = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, **kwargs, @@ -1523,6 +1556,7 @@ def chunk_kda_with_fused_gate( scale=scale, initial_state=initial_state.contiguous() if initial_state is not None else None, output_final_state=output_final_state, + lower_bound=lower_bound, cu_seqlens=cu_seqlens, ) return o, final_state @@ -1543,6 +1577,7 @@ def kda_gate_fwd_kernel( A, y, g_bias, + lower_bound, beta: tl.constexpr, threshold: tl.constexpr, T, @@ -1551,12 +1586,12 @@ def kda_gate_fwd_kernel( BT: tl.constexpr, BD: tl.constexpr, HAS_BIAS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, ): i_t, i_h = tl.program_id(0), tl.program_id(1) n_t = i_t * BT - b_a = tl.load(A + i_h).to(tl.float32) - b_a = -tl.exp(b_a) + b_a = tl.exp(tl.load(A + i_h).to(tl.float32)) stride_row = H * D stride_col = 1 @@ -1589,13 +1624,13 @@ def kda_gate_fwd_kernel( ) b_g = b_g + b_bias[None, :] - # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) - # When beta * x > threshold, use linear approximation x - # Use threshold to switch to linear when beta*x > threshold - g_scaled = b_g * beta - use_linear = g_scaled > threshold - sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) - b_y = b_a * sp + if USE_LOWER_BOUND: + b_y = lower_bound * tl.sigmoid(b_a * b_g) + else: + g_scaled = b_g * beta + use_linear = g_scaled > threshold + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) + b_y = -b_a * sp tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) @@ -1607,6 +1642,7 @@ def fused_kda_gate( g_bias: torch.Tensor | None = None, beta: float = 1.0, threshold: float = 20.0, + lower_bound: float | None = None, ) -> torch.Tensor: """ Forward pass for KDA gate: @@ -1634,6 +1670,7 @@ def fused_kda_gate( A, y, g_bias, + lower_bound or 0.0, beta, threshold, T, @@ -1641,6 +1678,7 @@ def fused_kda_gate( head_k_dim, BD=next_power_of_2(head_k_dim), HAS_BIAS=g_bias is not None, + USE_LOWER_BOUND=lower_bound is not None, ) y = y.view(*orig_shape, H, head_k_dim) diff --git a/vllm/transformers_utils/configs/kimi_k3.py b/vllm/transformers_utils/configs/kimi_k3.py new file mode 100644 index 00000000000..62de6dc03b4 --- /dev/null +++ b/vllm/transformers_utils/configs/kimi_k3.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multimodal configuration.""" + +from transformers.configuration_utils import PretrainedConfig + +from vllm.logger import init_logger +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig + +logger = init_logger(__name__) + + +class KimiK3VisionConfig(PretrainedConfig): + model_type = "kimi_k3_vision" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple[int, int] = (2, 2), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + mm_projector_type: str = "patchmergerv2", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + qkv_hidden_size: int = 1536, + norm_type: str = "rmsnorm", + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = "mlp2", + linear_bias: bool = False, + activation_func: str = "gelu_pytorch_tanh", + pos_emb_interpolation_mode: str = "bilinear", + text_hidden_size: int = 2304, + **kwargs, + ): + super().__init__(**kwargs) + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = tuple(merge_kernel_size) + self.video_attn_type = video_attn_type + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = ( + mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + ) + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + # Aliases consumed by the Kimi-K2.5 vision implementation. + self.num_attention_heads = vt_num_attention_heads + self.num_hidden_layers = vt_num_hidden_layers + self.hidden_size = vt_hidden_size + self.intermediate_size = vt_intermediate_size + + +class KimiK3Config(PretrainedConfig): + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig | None = None, + vision_config: dict | KimiK3VisionConfig | None = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + image_placeholder: str = "<|kimi_image_placeholder|>", + **kwargs, + ): + if text_config is None: + self.text_config = KimiLinearConfig() + elif isinstance(text_config, dict): + self.text_config = KimiLinearConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = KimiK3VisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = KimiK3VisionConfig(**vision_config) + else: + self.vision_config = vision_config + + # K3's vision projector output must match the text model's hidden size. + # Override any value provided in vision_config for safety. + if self.vision_config.text_hidden_size != self.text_config.hidden_size: + logger.info( + "Overriding vision_config.text_hidden_size from %s to %s " + "to match text_config.hidden_size", + self.vision_config.text_hidden_size, + self.text_config.hidden_size, + ) + self.vision_config.text_hidden_size = self.text_config.hidden_size + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.image_placeholder = image_placeholder + + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + @property + def vocab_size(self) -> int: + return self.text_config.vocab_size diff --git a/vllm/transformers_utils/processors/kimi_k3.py b/vllm/transformers_utils/processors/kimi_k3.py new file mode 100644 index 00000000000..73af85eb090 --- /dev/null +++ b/vllm/transformers_utils/processors/kimi_k3.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import BaseImageProcessor, BatchFeature, TensorType +from transformers.processing_utils import ProcessorMixin + +from vllm.tokenizers.hf import HfTokenizer + + +class KimiK3Processor(ProcessorMixin): + """HF-style processor wrapper for the image-only Kimi-K3 model. + + K3 exposes the standard ``image`` modality, so vLLM calls this processor + with ``images=[PIL, ...]``. The underlying checkpoint image processor + (``KimiK3VisionProcessor``) works on ``{"type": "image", "image": PIL}`` + media dicts, so this wrapper adapts bare PIL images into that shape before + delegating to ``preprocess``. + + Text is only tokenized here; the single ``<|kimi_image_placeholder|>`` + token per image is expanded into the resolution-aware media block by the + model's ``_get_prompt_updates`` on the vLLM side. + """ + + attributes = ["image_processor", "tokenizer"] + + def __init__( + self, + image_processor: BaseImageProcessor, + tokenizer: HfTokenizer, + ) -> None: + self.image_processor = image_processor + self.tokenizer = tokenizer + + def __call__( + self, + text: str | list[str] | None = None, + images: object | list[object] | None = None, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + if images is not None: + if not isinstance(images, list): + images = [images] + medias = [{"type": "image", "image": image} for image in images] + mm_inputs = self.image_processor.preprocess( + medias, + return_tensors=return_tensors, + ) + else: + mm_inputs = {} + + if text is not None: + if not isinstance(text, list): + text = [text] + text_inputs = self.tokenizer(text) + else: + text_inputs = {} + + return BatchFeature( + data={**text_inputs, **mm_inputs}, + tensor_type=return_tensors, + ) diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py index e8f90efce98..76dd1982202 100644 --- a/vllm/v1/attention/ops/triton_merge_attn_states.py +++ b/vllm/v1/attention/ops/triton_merge_attn_states.py @@ -142,6 +142,15 @@ def merge_attn_states( # attention backend. prefix_head_stride = prefix_output.stride(1) output_head_stride = output.stride(1) + # lse tensors are [NUM_HEADS, NUM_TOKENS] but may be non-contiguous views + # (e.g. a transpose of a backend's [NUM_TOKENS, NUM_HEADS] output), so index + # them by their actual strides rather than assuming a contiguous layout. + prefix_lse_head_stride = prefix_lse.stride(0) + prefix_lse_token_stride = prefix_lse.stride(1) + suffix_lse_head_stride = suffix_lse.stride(0) + suffix_lse_token_stride = suffix_lse.stride(1) + output_lse_head_stride = output_lse.stride(0) if output_lse is not None else 0 + output_lse_token_stride = output_lse.stride(1) if output_lse is not None else 0 # If prefill_tokens_with_context is None, all tokens should use prefix context if prefill_tokens_with_context is None: @@ -157,6 +166,12 @@ def merge_attn_states( suffix_lse, prefix_head_stride, output_head_stride, + prefix_lse_head_stride, + prefix_lse_token_stride, + suffix_lse_head_stride, + suffix_lse_token_stride, + output_lse_head_stride, + output_lse_token_stride, output_scale, prefill_tokens_with_context, head_size, @@ -176,6 +191,12 @@ def merge_attn_states_kernel( suffix_lse, # [NUM_HEADS, NUM_TOKENS] prefix_head_stride, output_head_stride, + prefix_lse_head_stride, + prefix_lse_token_stride, + suffix_lse_head_stride, + suffix_lse_token_stride, + output_lse_head_stride, + output_lse_token_stride, output_scale, # scale tensor or None prefill_tokens_with_context, HEAD_SIZE: tl.constexpr, @@ -186,7 +207,6 @@ def merge_attn_states_kernel( FP8_MAX: tl.constexpr = float8_info.max, ): token_idx = tl.program_id(0) - num_tokens = tl.num_programs(0) head_idx = tl.program_id(1) num_heads = tl.num_programs(1) @@ -198,9 +218,18 @@ def merge_attn_states_kernel( # For tokens without context (token_idx >= prefill_tokens_with_context), # directly copy from suffix_output if not prefix_mask: - s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx) + s_lse = tl.load( + suffix_lse + + head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride + ) if OUTPUT_LSE: - tl.store(output_lse + head_idx * num_tokens + token_idx, s_lse) + tl.store( + output_lse + + head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride, + s_lse, + ) s_out = tl.load( suffix_output @@ -227,8 +256,16 @@ def merge_attn_states_kernel( # For tokens with context (token_idx < prefill_tokens_with_context), # perform normal merge operation - p_lse = tl.load(prefix_lse + head_idx * num_tokens + token_idx) - s_lse = tl.load(suffix_lse + head_idx * num_tokens + token_idx) + p_lse = tl.load( + prefix_lse + + head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride + ) + s_lse = tl.load( + suffix_lse + + head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride + ) # FA2 and FA3 have different behavior for when the sum-exp is 0, this namely # arises with 0 len seqlens. FA3 returns -inf here while FA2 returns inf. @@ -251,7 +288,12 @@ def merge_attn_states_kernel( # Both sides empty (max_lse == -inf) => undefined merge; keep -inf so # downstream merges continue to treat the token as empty. out_lse = tl.where(max_lse == float("-inf"), float("-inf"), out_lse) - tl.store(output_lse + head_idx * num_tokens + token_idx, out_lse) + tl.store( + output_lse + + head_idx * output_lse_head_stride + + token_idx * output_lse_token_stride, + out_lse, + ) p_out = tl.load( prefix_output From 7de49bab7e91a2a77e98a31fae4b1c615b34dce8 Mon Sep 17 00:00:00 2001 From: zofia <110436990+zufangzhu@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:16:16 +0800 Subject: [PATCH 67/67] [XPU][UT][CI] add xpu config to run gpt-oss accuracy in ut and ci (#48703) Signed-off-by: Zhu, Zufang Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/intel_jobs/test-intel.yaml | 24 +++++++++++++++++++ .../configs/gpt-oss-20b-xpu-baseline.yaml | 5 ++++ .../configs/gpt-oss-20b-xpu-triton-attn.yaml | 6 +++++ tests/evals/gpt_oss/configs/models-xpu.txt | 3 +++ 4 files changed, 38 insertions(+) create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml create mode 100644 tests/evals/gpt_oss/configs/models-xpu.txt diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 3fadb07f391..ec5cb2fd9e7 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -147,6 +147,30 @@ steps: 'cd tests && pytest -v -s quantization/test_auto_round.py && pytest -v -s quantization/test_online.py' + - label: "XPU GPQA Eval (GPT-OSS)" + depends_on: + - image-build-xpu + timeout_in_minutes: 60 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/ + - tests/evals/gpt_oss/ + - .buildkite/intel_jobs/test-intel.yaml + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'pip install "gpt-oss[eval]==0.0.5" && + cd tests && + pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' - label: "XPU compressed tensors FP8 test" depends_on: - image-build-xpu diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml new file mode 100644 index 00000000000..78a583888ba --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-baseline.yaml @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml new file mode 100644 index 00000000000..e711ffb331e --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-xpu-triton-attn.yaml @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: openai/gpt-oss-20b +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend TRITON_ATTN" diff --git a/tests/evals/gpt_oss/configs/models-xpu.txt b/tests/evals/gpt_oss/configs/models-xpu.txt new file mode 100644 index 00000000000..a9de9266b14 --- /dev/null +++ b/tests/evals/gpt_oss/configs/models-xpu.txt @@ -0,0 +1,3 @@ +# Intel XPU model configurations for GPQA evaluation +gpt-oss-20b-xpu-baseline.yaml +gpt-oss-20b-xpu-triton-attn.yaml