diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 8c2f98d8fe..15a7d0cde6 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # NOTE: This is currently targeting nv-dev branch due to sm120 support - set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # TODO: switch to nv_dev branch after it support situ + set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0836f11e26..e434c57012 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -547,6 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ | | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A+ | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I+ | `moonshotai/Kimi-K2.5` | | ✅︎ | +| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I+ | `moonshotai/Kimi-K3` | | ✅︎ | | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I+ | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ | | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I+ | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ | | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I+ | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ | diff --git a/pyproject.toml b/pyproject.toml index a4c7ff487a..0766645fc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,6 +155,8 @@ view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" [tool.typos.default.extend-words] +Hel = "Hel" +wether = "wether" iy = "iy" indx = "indx" # intel cpu features diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 9d93b60320..2a8b2734f8 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -16,8 +16,8 @@ PyNvVideoCodec==2.0.4 --extra-index-url https://flashinfer.ai/whl/ flashinfer-python==0.6.15.post1 flashinfer-cubin==0.6.15.post1 -apache-tvm-ffi==0.1.10 -tilelang==0.1.9 +apache-tvm-ffi==0.1.11 +tilelang==0.1.12 nvidia-cudnn-frontend>=1.19.1 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1 nvtx==0.2.15 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 6c155d89bd..4d1734e15f 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -45,7 +45,7 @@ anyio==4.14.1 # sse-starlette # starlette # watchfiles -apache-tvm-ffi==0.1.10 +apache-tvm-ffi==0.1.11 # via # -c requirements/cuda.txt # xgrammar diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index a9c6e9f5b1..298fb23f1d 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantDesc, fp8_w8a8_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( + deepgemm_moe_permute, +) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) @@ -41,6 +44,35 @@ from vllm.utils.deep_gemm import ( BLOCK_SIZE = [128, 128] +@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") +def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init): + hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16) + activations, scales = per_token_group_quant_fp8( + hidden_states, + group_size=128, + use_ue8m0=True, + ) + topk_ids = torch.tensor([[0], [1]], device="cuda", dtype=torch.int64) + + _, permuted_scales, expert_ids, _, _ = deepgemm_moe_permute( + aq=activations, + aq_scale=scales, + topk_ids=topk_ids, + local_num_experts=2, + expert_map=None, + expert_tokens_meta=None, + ) + + padding = expert_ids < 0 + assert padding.any() + torch.testing.assert_close( + permuted_scales[padding], + torch.zeros_like(permuted_scales[padding]), + rtol=0, + atol=0, + ) + + def make_block_quant_fp8_weights( e: int, n: int, diff --git a/tests/models/registry.py b/tests/models/registry.py index 7ed8d6225f..24bc4a905a 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1026,6 +1026,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { "moonshotai/Kimi-K2.5", trust_remote_code=True, ), + "KimiK3ForConditionalGeneration": _HfExamplesInfo( + "moonshotai/Kimi-K3", + trust_remote_code=True, + ), "KimiVLForConditionalGeneration": _HfExamplesInfo( "moonshotai/Kimi-VL-A3B-Instruct", extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, @@ -1466,6 +1470,14 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { is_available_online=False, use_original_num_layers=True, # DSpark has >1 draft block ), + "K3DSparkModel": _HfExamplesInfo( + "moonshotai/Kimi-K3", + speculative_model="Inferact/Kimi-K3-DSpark", + use_original_num_layers=True, # DSpark has >1 draft block + trust_remote_code=True, + # FIXME: Investigate the NVML failure in CI. + is_available_online=False, + ), "Qwen3DSparkModel": _HfExamplesInfo( "Qwen/Qwen3-8B", speculative_model="deepseek-ai/dspark_qwen3_8b_block7", @@ -1664,6 +1676,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { trust_remote_code=True, max_model_len=4096, ), + "KimiK3MTPModel": _HfExamplesInfo( + "moonshotai/Kimi-K3", + speculative_model="moonshotai/Kimi-K3", + trust_remote_code=True, + is_available_online=False, + ), "LongCatFlashMTPModel": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True, diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py new file mode 100644 index 0000000000..28dcd8ebf4 --- /dev/null +++ b/tests/models/test_dspark_mla.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.registry import ModelRegistry +from vllm.models.kimi_k3.nvidia import dspark_mla +from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel + + +def test_dspark_mla_uses_compile_free_model_entrypoint(): + assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM + assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper) + + +@pytest.mark.parametrize( + ("checkpoint_name", "runtime_name", "shard_id"), + [ + ( + "layers.0.self_attn.q_a_proj.weight", + "model.layers.0.self_attn.fused_qkv_a_proj.weight", + 0, + ), + ( + "layers.0.self_attn.kv_a_proj_with_mqa.weight", + "model.layers.0.self_attn.fused_qkv_a_proj.weight", + 1, + ), + ( + "layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 0, + ), + ( + "layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + 1, + ), + ("context_proj.weight", "model.context_proj.weight", None), + ], +) +def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id): + assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard( + checkpoint_name + ) == (runtime_name, shard_id) + + +def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head(): + assert not K3DSparkForCausalLM.has_own_embed_tokens + assert not K3DSparkForCausalLM.has_own_lm_head + assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == { + "confidence_head", + "embed_tokens", + "lm_head", + } + + +@pytest.mark.cpu_test +def test_dspark_markov_head_is_replicated( + monkeypatch: pytest.MonkeyPatch, +): + from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding + + monkeypatch.setattr( + vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3 + ) + monkeypatch.setattr( + vocab_parallel_embedding, + "get_tensor_model_parallel_world_size", + lambda: 8, + ) + monkeypatch.setattr( + logits_processor, + "get_current_vllm_config", + lambda: SimpleNamespace(model_config=None), + ) + + head = DSparkMarkovHead(128, 128, 8, prefix="markov_head") + assert head.markov_w2.tp_size == 1 + assert head.markov_w1.weight.shape == (128, 8) + assert head.markov_w2.weight.shape == (128, 8) + + def fail_collective(*args, **kwargs): + raise AssertionError("replicated Markov head must not invoke TP collectives") + + monkeypatch.setattr( + vocab_parallel_embedding, + "tensor_model_parallel_all_reduce", + fail_collective, + ) + logits_processor = LogitsProcessor(128) + monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective) + + markov_embed = head.embed(torch.tensor([1, 2])) + bias = head.bias(markov_embed, logits_processor) + assert markov_embed.shape == (2, 8) + assert bias.shape == (2, 128) + + +@pytest.mark.cpu_test +def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): + markov_head_calls = [] + + class DummyModule(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def make_markov_head(*args, **kwargs): + markov_head_calls.append((args, kwargs)) + return DummyModule() + + monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None) + monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule) + monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule) + monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule) + monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head) + + config = SimpleNamespace( + target_hidden_size=16, + num_target_layers=2, + hidden_size=8, + rms_norm_eps=1e-6, + num_hidden_layers=1, + vocab_size=128, + draft_vocab_size=128, + markov_rank=4, + ) + vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=config) + ) + ) + + K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model") + + assert len(markov_head_calls) == 1 diff --git a/tests/test_config.py b/tests/test_config.py index 8171950cdd..21bdde0457 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1235,6 +1235,26 @@ def test_vllm_config_defaults_are_none(): assert getattr(config.compilation_config, k) is None +def test_validate_mamba_align_subblock_prefill(): + """Align mode permits configured prefill chunks smaller than a block.""" + config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=11392, + mamba_cache_mode="align", + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8192, + long_prefill_token_threshold=4096, + disable_chunked_mm_input=False, + ), + ) + + VllmConfig.validate_block_size(config) + + @pytest.mark.parametrize( ("model_id", "compilation_config", "optimization_level"), [ diff --git a/tests/transformers_utils/test_dspark_mla_config.py b/tests/transformers_utils/test_dspark_mla_config.py new file mode 100644 index 0000000000..f43d668b73 --- /dev/null +++ b/tests/transformers_utils/test_dspark_mla_config.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig +from vllm.transformers_utils.config import get_config +from vllm.transformers_utils.configs.k3_dspark import K3DSparkConfig + + +def _write_dspark_config(path, **overrides): + path.mkdir() + config = { + "architectures": ["K3DSparkModel"], + "model_type": "k3_dspark", + "hidden_size": 7168, + "intermediate_size": 14336, + "num_hidden_layers": 5, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "q_lora_rank": 1536, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "vocab_size": 163840, + "rms_norm_eps": 1e-5, + "max_position_embeddings": 32768, + "rope_theta": 50000.0, + "num_target_layers": 5, + "target_hidden_size": 7168, + "target_num_hidden_layers": 93, + "target_layer_ids": [2, 23, 47, 71, 89], + "markov_rank": 256, + "draft_vocab_size": 163840, + "torch_dtype": "bfloat16", + } + config.update(overrides) + (path / "config.json").write_text(json.dumps(config)) + + +def _write_target_config(path): + path.mkdir() + config = { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "hidden_size": 7168, + "intermediate_size": 14336, + "num_hidden_layers": 93, + "num_attention_heads": 56, + "num_key_value_heads": 8, + "vocab_size": 163840, + "max_position_embeddings": 32768, + "torch_dtype": "bfloat16", + } + (path / "config.json").write_text(json.dumps(config)) + + +def test_dspark_mla_config_loads_from_local_json(tmp_path): + draft_path = tmp_path / "draft" + _write_dspark_config(draft_path) + + config = get_config(draft_path, trust_remote_code=False) + + assert isinstance(config, K3DSparkConfig) + assert config.model_type == "k3_dspark" + assert config.architectures == ["K3DSparkModel"] + assert config.hidden_act == "silu" + assert config.rope_parameters == { + "rope_type": "default", + "rope_theta": 50000.0, + } + assert config.n_routed_experts == 0 + assert config.draft_vocab_size == config.vocab_size + + +@pytest.mark.parametrize( + "overrides", + [ + {"mla_use_nope": True}, + {"mla_use_output_gate": True}, + {"mla_use_qk_norm": True}, + {"dspark_bonus_anchor": True}, + {"q_lora_rank": None}, + {"draft_vocab_size": 8192}, + {"target_layer_ids": []}, + {"num_target_layers": 4}, + ], +) +def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides): + draft_path = tmp_path / "draft" + _write_dspark_config(draft_path, **overrides) + + with pytest.raises(ValueError, match="MLA DSpark"): + get_config(draft_path, trust_remote_code=False) + + +def test_dspark_mla_uses_latent_kv_geometry(tmp_path): + draft_path = tmp_path / "draft" + _write_dspark_config(draft_path) + model_config = ModelConfig( + model=str(draft_path), + tokenizer_mode="skip", + runner="draft", + max_model_len=32768, + ) + + assert model_config.is_deepseek_mla + assert model_config.use_mla + assert model_config.get_head_size() == 576 + # external_launcher skips ParallelConfig's local-GPU-count check so the + # config logic can be exercised at TP8 on a single-GPU test node. + parallel_config = ParallelConfig( + tensor_parallel_size=8, distributed_executor_backend="external_launcher" + ) + assert model_config.get_num_kv_heads(parallel_config) == 1 + assert model_config.get_num_attention_heads(parallel_config) == 8 + assert model_config.get_num_experts() == 0 + + +def test_dspark_mla_speculative_config_preserves_architecture(tmp_path): + target_path = tmp_path / "target" + draft_path = tmp_path / "draft" + _write_target_config(target_path) + _write_dspark_config(draft_path) + target_config = ModelConfig( + model=str(target_path), tokenizer_mode="skip", max_model_len=32768 + ) + speculative_config = SpeculativeConfig( + model=str(draft_path), + method="dspark", + num_speculative_tokens=8, + target_model_config=target_config, + target_parallel_config=ParallelConfig(), + ) + + assert speculative_config.parallel_drafting + assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"] + assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark" + assert speculative_config.draft_model_config.use_mla + + +def test_dspark_mla_rejects_decode_context_parallelism(tmp_path): + target_path = tmp_path / "target" + draft_path = tmp_path / "draft" + _write_target_config(target_path) + _write_dspark_config(draft_path) + target_config = ModelConfig( + model=str(target_path), tokenizer_mode="skip", max_model_len=32768 + ) + + with pytest.raises(ValueError, match="does not currently support decode context"): + SpeculativeConfig( + model=str(draft_path), + method="dspark", + num_speculative_tokens=8, + target_model_config=target_config, + target_parallel_config=ParallelConfig( + tensor_parallel_size=2, + decode_context_parallel_size=2, + distributed_executor_backend="external_launcher", + ), + ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index c520f22e7e..36ebd1af8f 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -73,7 +73,11 @@ DEVICE_TYPE = current_platform.device_type def test_mla_kv_cache_spec_uses_layer_cache_dtype( cache_dtype: str, expected_quant_mode: KVQuantMode ): - layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576) + layer = SimpleNamespace( + kv_cache_dtype=cache_dtype, + head_size=576, + non_causal_multi_token_decode=False, + ) vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=64), model_config=None ) diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py new file mode 100644 index 0000000000..7aca8400f9 --- /dev/null +++ b/tests/v1/attention/test_mla_noncausal.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonMetadata, + MLACommonMetadataBuilder, + QueryLenSupport, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import MLAAttentionSpec + + +class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + supports_non_causal_multi_token_decode = True + + +def _metadata( + query_start_loc: list[int], num_tokens: int | None = None +) -> CommonAttentionMetadata: + num_reqs = len(query_start_loc) - 1 + num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens + return CommonAttentionMetadata( + query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), + query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), + seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8, + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=max( + end - start for start, end in zip(query_start_loc, query_start_loc[1:]) + ), + max_seq_len=num_reqs * 100 + 8, + block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view( + num_reqs, 3 + ), + slot_mapping=torch.arange(num_tokens), + causal=False, + seq_lens_cpu_upper_bound=None, + ) + + +def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: + builder = object.__new__(_NonCausalMLAMetadataBuilder) + builder.device = torch.device("cpu") + builder.reorder_batch_threshold = 1 + builder.query_len_support = QueryLenSupport.SINGLE_ONLY + builder.non_causal_multi_token_decode = marked + builder.dcp_world_size = 1 + builder.metadata_cls = MLACommonMetadata + builder.model_config = SimpleNamespace( + dtype=torch.bfloat16, get_head_size=lambda: 576 + ) + return builder + + +def test_noncausal_block_uses_decode_without_cpu_lengths(): + common_metadata = _metadata([0, 8, 16]) + metadata = _builder().build(0, common_metadata) + + assert metadata.num_decodes == 2 + assert metadata.num_decode_tokens == 16 + assert metadata.num_prefills == 0 + assert metadata.prefill is None + assert metadata.decode is not None + assert metadata.decode.block_table.shape == (2, 3) + assert metadata.decode.seq_lens.shape == (2,) + assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor) + assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens) + assert not metadata.causal + + +def test_noncausal_support_is_explicit_and_uniform(): + with pytest.raises(ValueError, match="explicitly supported"): + _builder(marked=False).build(0, _metadata([0, 8, 16])) + with pytest.raises(ValueError, match="uniform query block"): + _builder().build(0, _metadata([0, 3, 8])) + + +def test_noncausal_block_allows_trailing_cudagraph_padding(): + common_metadata = _metadata([0, 8, 16, 16], num_tokens=24) + common_metadata.seq_lens[-1] = 0 + + metadata = _builder().build(0, common_metadata) + + assert metadata.num_decodes == 3 + assert metadata.num_decode_tokens == 24 + assert metadata.decode is not None + assert metadata.decode.seq_lens.tolist() == [108, 208, 0] + + +def test_noncausal_block_rejects_non_trailing_padding(): + with pytest.raises(ValueError, match="uniform query block"): + _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24)) + + +def test_noncausal_decode_metadata_keeps_live_request_buffers(): + common_metadata = _metadata([0, 8, 16]) + metadata = _builder().build(0, common_metadata) + + assert metadata.decode is not None + assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr() + assert ( + metadata.decode.block_table.data_ptr() + == common_metadata.block_table_tensor.data_ptr() + ) + + +def test_mla_cache_marker_is_promoted_to_group_capability(): + kwargs = { + "block_size": 64, + "num_kv_heads": 1, + "head_size": 576, + "dtype": torch.bfloat16, + } + marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True) + unmarked = MLAAttentionSpec(**kwargs) + + assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode + assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode + assert not MLAAttentionSpec.merge( + [unmarked, unmarked] + ).non_causal_multi_token_decode diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 70457d5af9..ec3f293c0a 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -40,6 +40,8 @@ def test_mamba_align_split_partial_tail_schedule(): hash_block_size = 32 mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=8192, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=hash_block_size, mamba_partial_cache_hit=True, @@ -75,6 +77,78 @@ def test_mamba_align_split_partial_tail_schedule(): assert split(self=mock, request=req2, num_new_tokens=1000) == 512 +def test_mamba_align_split_when_block_exceeds_scheduling_budget(): + """Sub-block chunks make progress only when no step can fit a full block.""" + block_size = 11392 + token_budget = 8192 + prompt_length = 30000 + mock = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=token_budget, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + use_eagle=False, + hash_block_size=32, + mamba_partial_cache_hit=False, + ) + req = make_request("0", [0] * prompt_length, 32, sha256) + split = Scheduler._mamba_block_aligned_split + + mock.max_num_scheduled_tokens = block_size + assert split(self=mock, request=req, num_new_tokens=token_budget) == 0 + mock.max_num_scheduled_tokens = token_budget + + scheduled_chunks = [] + while req.num_computed_tokens < prompt_length: + num_new_tokens = min(token_budget, prompt_length - req.num_computed_tokens) + num_scheduled_tokens = split( + self=mock, + request=req, + num_new_tokens=num_new_tokens, + ) + assert 0 < num_scheduled_tokens <= token_budget + scheduled_chunks.append(num_scheduled_tokens) + req.num_computed_tokens += num_scheduled_tokens + + assert scheduled_chunks == [8192, 3200, 8192, 3200, 7216] + + +def test_mamba_align_split_when_block_exceeds_long_prefill_threshold(): + """A long-prefill cap below the block size permits sub-block progress.""" + block_size = 512 + token_budget = 8192 + long_prefill_threshold = 384 + prompt_length = 1300 + mock = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=token_budget, + scheduler_config=SimpleNamespace( + long_prefill_token_threshold=long_prefill_threshold + ), + use_eagle=False, + hash_block_size=32, + mamba_partial_cache_hit=False, + ) + req = make_request("0", [0] * prompt_length, 32, sha256) + split = Scheduler._mamba_block_aligned_split + + scheduled_chunks = [] + while req.num_computed_tokens < prompt_length: + num_new_tokens = min( + long_prefill_threshold, + prompt_length - req.num_computed_tokens, + ) + num_scheduled_tokens = split( + self=mock, + request=req, + num_new_tokens=num_new_tokens, + ) + assert 0 < num_scheduled_tokens <= long_prefill_threshold + scheduled_chunks.append(num_scheduled_tokens) + req.num_computed_tokens += num_scheduled_tokens + + assert scheduled_chunks == [384, 128, 384, 128, 276] + + def test_hybrid_mamba_align_partial_hash_hit(): hash_block_size = 2 mamba_block_size = 2 * hash_block_size diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index dd89769315..13ed7c7b9d 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1078,6 +1078,8 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): # Create minimal mock with just the needed attributes mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), + max_num_scheduled_tokens=3 * block_size, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=block_size, mamba_partial_cache_hit=False, diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index e8a86b10e8..dc9c1ffb30 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -333,6 +333,54 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) +def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): + """Handle an eager attention break inside a multi-stream overlap region.""" + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) + from vllm.utils.multi_stream_utils import maybe_execute_in_parallel + + x = torch.zeros(1024, device="cuda") + output = torch.empty_like(x) + aux_stream = torch.cuda.Stream() + main_event = torch.cuda.Event() + aux_event = torch.cuda.Event() + + @eager_break_during_capture + def attention(query: torch.Tensor, out: torch.Tensor) -> None: + torch.add(query, 3.0, out=out) + + def attention_frontend() -> torch.Tensor: + query = x * 2.0 + attention_output = torch.empty_like(x) + attention(query, attention_output) + return attention_output + + cap = BreakableCUDAGraphCapture() + with cap: + attention_output, gate = maybe_execute_in_parallel( + attention_frontend, + lambda: x * 5.0, + main_event, + aux_event, + aux_stream, + ) + torch.add(attention_output, gate, out=output) + + assert cap.num_graphs == 2 + assert cap.num_eager_breaks == 1 + + for value in (1.0, 7.0, 19.0): + x.fill_(value) + cap.replay() + cuda_capture_stream.synchronize() + torch.testing.assert_close( + output, + torch.full_like(output, value * 7.0 + 3.0), + ) + + # --------------------------------------------------------------------------- # Replay ordering # --------------------------------------------------------------------------- diff --git a/tests/v1/spec_decode/test_llm_base_proposer_sampling.py b/tests/v1/spec_decode/test_llm_base_proposer_sampling.py index 9c7ec760eb..cb390b992f 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer_sampling.py +++ b/tests/v1/spec_decode/test_llm_base_proposer_sampling.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest import torch from vllm.platforms import current_platform @@ -8,6 +11,7 @@ from vllm.utils.torch_utils import set_random_seed from vllm.v1.sample.logits_processor import LogitsProcessors from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.llm_base_proposer import ( + SpecDecodeBaseProposer, compute_probs_and_sample_next_token, ) @@ -68,3 +72,21 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race(): assert torch.equal(actual_ids, expected_ids) assert torch.allclose(actual_probs, probs) + + +@pytest.mark.parametrize( + ("architecture", "expected"), + [ + ("DeepSeekMTPModel", True), + ("KimiK3MTPModel", True), + ("MiniMaxM3ForCausalLM", False), + ], +) +def test_mtp_model_returns_tuple(architecture: str, expected: bool): + proposer = object.__new__(SpecDecodeBaseProposer) + proposer.method = "mtp" + proposer.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=[architecture]) + ) + + assert proposer.model_returns_tuple() is expected diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 52f73e9050..1d6ff5a80b 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,9 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This is currently targeting nv-dev branch due to sm120 support -DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a" +DEEPGEMM_GIT_REF="f5a76426fa084087169693fd0cd815223576d6e9" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9ab6c81dce..2a6bdeaa2a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -133,6 +133,8 @@ def _rocm_aiter_fused_moe_impl( bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe @@ -143,6 +145,12 @@ def _rocm_aiter_fused_moe_impl( extra_kwargs: dict = {} if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): extra_kwargs["gate_mode"] = gate_mode + if ( + getattr(ActivationType, "Situv2", None) is not None + and activation == ActivationType.Situv2 + ): + extra_kwargs["beta"] = beta + extra_kwargs["linear_beta"] = linear_beta return fused_moe( hidden_states, @@ -193,6 +201,8 @@ def _rocm_aiter_fused_moe_fake( bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: if output_dtype is not None: return torch.empty_like(hidden_states, dtype=output_dtype) @@ -1618,6 +1628,7 @@ class rocm_aiter_ops: "silu": ActivationType.Silu, "gelu": ActivationType.Gelu, "swiglu": ActivationType.Swiglu, + "situ": getattr(ActivationType, "Situv2", None), } return mapping.get(name) @@ -2222,6 +2233,8 @@ class rocm_aiter_ops: bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, swiglu_limit: float = 0.0, + beta: float | None = None, + linear_beta: float | None = None, ) -> torch.Tensor: return torch.ops.vllm.rocm_aiter_fused_moe( hidden_states, @@ -2246,6 +2259,8 @@ class rocm_aiter_ops: bias2, moe_sorting_dispatch_policy, swiglu_limit, + beta, + linear_beta, ) @staticmethod diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index fbc5dcff40..4eb66613a8 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -771,7 +771,6 @@ class CompilationConfig: "vllm::qwen_gdn_attention_core", "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", - "vllm::kda_attention", "vllm::sparse_attn_indexer", "vllm::rocm_aiter_sparse_attn_indexer", "vllm::deepseek_v4_attention", diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 5ebc410f3c..2e908ead17 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -1028,11 +1028,6 @@ class ParallelConfig: "Disabled the custom all-reduce kernel because it is not " "supported on current platform." ) - if self.nnodes > 1: - self.disable_custom_all_reduce = True - logger.debug( - "Disabled the custom all-reduce since we are running on multi-node." - ) if self.ray_workers_use_nsight and not self.use_ray: raise ValueError( "Unable to use nsight profiling unless workers run with Ray." diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index e50fe4b5d9..26d091520b 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -51,6 +51,7 @@ MTPModelTypes = Literal[ "minimax_m3_mtp", "bailing_hybrid_mtp", "mtp", + "kimi_k3_mtp", "pangu_ultra_moe_mtp", "step3p5_mtp", "hy_v3_mtp", @@ -355,6 +356,16 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]} ) + if hf_config.model_type == "kimi_k3": + # Kimi-K3 keeps the text-model fields (incl. the MTP layer count) + # nested under ``text_config`` (a KimiLinearConfig). + text_config = getattr(hf_config, "text_config", hf_config) + n_predict = getattr(text_config, "num_nextn_predict_layers", None) + hf_config.model_type = "kimi_k3_mtp" + hf_config.update( + {"n_predict": n_predict, "architectures": ["KimiK3MTPModel"]} + ) + if hf_config.architectures[0] == "MiMoForCausalLM": hf_config.model_type = "mimo_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", None) @@ -943,6 +954,7 @@ class SpeculativeConfig: if self.method == "dspark" and ( "Qwen3DSparkModel" not in self.draft_model_config.architectures and "Gemma4DSparkModel" not in self.draft_model_config.architectures + and "K3DSparkModel" not in self.draft_model_config.architectures ): # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config # and its weights ship in the target checkpoint. @@ -972,6 +984,16 @@ class SpeculativeConfig: if self.method in ("dflash", "dspark"): self.parallel_drafting = True + if ( + self.method == "dspark" + and "K3DSparkModel" in self.draft_model_config.architectures + and self.target_parallel_config.decode_context_parallel_size > 1 + ): + raise ValueError( + "MLA DSpark does not currently support decode context " + "parallelism; set decode_context_parallel_size=1." + ) + if self.num_speculative_tokens is not None and hasattr( self.draft_model_config.hf_config, "num_lookahead_tokens" ): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 82ca8ea26a..dce048dd31 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -72,11 +72,34 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( "GraniteMoeForCausalLM", "InklingForCausalLM", "InklingForConditionalGeneration", + "KimiK3ForConditionalGeneration", "LongcatFlashNgramForCausalLM", "Qwen2MoeForCausalLM", } ) +# Architectures that default to V1 on ROCm: the V2 runner faults during the +# profile run. VLLM_USE_V2_MODEL_RUNNER=1 still forces V2. +# TODO: fix V2 enablement +ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "KimiK3ForConditionalGeneration", + } +) + + +@lru_cache +def default_v2_model_runner_architectures() -> frozenset[str]: + """Architectures defaulting to the V2 model runner on this platform.""" + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + return ( + DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES + - ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES + ) + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -619,16 +642,20 @@ class VllmConfig: if model_config.runner_type != "generate": return False - if getattr(model_config, "is_hybrid", False): + architectures = getattr(model_config, "architectures", []) + default_architectures = default_v2_model_runner_architectures() + is_default_v2_architecture = any( + arch in default_architectures for arch in architectures + ) + + if getattr(model_config, "is_hybrid", False) and ( + not is_default_v2_architecture + ): return False if getattr(model_config, "is_attention_free", False): return False - architectures = getattr(model_config, "architectures", []) - return ( - any(arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures) - or not model_config.is_moe - ) + return is_default_v2_architecture or not model_config.is_moe @property def needs_dp_coordinator(self) -> bool: @@ -1196,6 +1223,9 @@ class VllmConfig: "DeepSeekV4MTPModel", "InklingForCausalLM", "InklingForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "KimiK3MTPModel", + "KimiLinearForCausalLM", "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", ) @@ -2277,14 +2307,6 @@ class VllmConfig: # Mamba cache align-mode constraints if self.cache_config.mamba_cache_mode == "align": - assert block_size <= self.scheduler_config.max_num_batched_tokens, ( - "In Mamba cache align mode, block_size " - f"({block_size}) must be <= " - "max_num_batched_tokens " - f"({self.scheduler_config.max_num_batched_tokens})." - ) - if self.scheduler_config.long_prefill_token_threshold > 0: - assert self.scheduler_config.long_prefill_token_threshold >= block_size assert not self.scheduler_config.disable_chunked_mm_input, ( "Chunked MM input is required because we need the flexibility " "to schedule a multiple of block_size tokens even if they are " diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 23e37ca830..06b441c5a4 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -340,6 +340,18 @@ class CudaCommunicator(DeviceCommunicatorBase): torch.distributed.all_reduce(out, group=self.device_group) return out + def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None: + ca_comm = self.ca_comm + if ca_comm is None: + return None + return ca_comm.custom_all_gather(input_.contiguous()) + + def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None: + ca_comm = self.ca_comm + if ca_comm is None: + return None + return ca_comm.custom_reduce_scatter(input_.contiguous()) + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: # Route uniform dim-0 all-gathers through NVLS symmetric memory when # enabled (mirrors reduce_scatter); otherwise fall back to the diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 95db6cc924..79a37c1805 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -25,6 +25,11 @@ except Exception: # For CPUs custom_ar = False +try: + import torch.distributed._symmetric_memory as torch_symm_mem +except ImportError: + torch_symm_mem = None + logger = init_logger(__name__) @@ -49,7 +54,17 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402 class CustomAllreduce: - _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8] + _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16] + _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024 + _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = { + 2: 8 * 1024 * 1024, + 4: 4 * 1024 * 1024, + 6: 2 * 1024 * 1024, + 8: 2 * 1024 * 1024, + 16: 2 * 1024 * 1024, + } + _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 + _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024 # max_size: max supported allreduce size def __init__( @@ -57,6 +72,10 @@ class CustomAllreduce: group: ProcessGroup, device: int | str | torch.device, max_size=8192 * 1024, + max_all_gather_size=_DEFAULT_ALL_GATHER_MAX_SIZE, + max_mnnvl_all_gather_size=None, + max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE, + max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE, symm_mem_enabled=False, ) -> None: """ @@ -70,7 +89,20 @@ class CustomAllreduce: are in the same node. """ self._IS_CAPTURING = False + self._ptr = 0 self.disabled = True + self.mnnvl_buffer = None + self.mnnvl_handle = None + self.mnnvl_peer_buffers: list[torch.Tensor] | None = None + self.mnnvl_multicast_ptr = 0 + self.mnnvl_buffer_size = 0 + self.mnnvl_lamport_ag_local_ptr = 0 + self.mnnvl_lamport_ag_multicast_ptr = 0 + self.mnnvl_lamport_rs_local_ptr = 0 + self.mnnvl_lamport_epochs = None + self.mnnvl_lamport_ag_epoch_ptr = 0 + self.mnnvl_lamport_rs_epoch_ptr = 0 + self.mnnvl_only = False if not custom_ar: # disable because of missing custom allreduce library @@ -87,13 +119,8 @@ class CustomAllreduce: "CustomAllreduce should be attached to a non-NCCL group." ) - if not all(in_the_same_node_as(group, source_rank=0)): - # No need to initialize custom allreduce for multi-node case. - logger.warning( - "Custom allreduce is disabled because this process group" - " spans across nodes." - ) - return + same_node = all(in_the_same_node_as(group, source_rank=0)) + self.mnnvl_only = not same_node rank = dist.get_rank(group=self.group) self.rank = rank @@ -126,28 +153,30 @@ class CustomAllreduce: and device_capability is not None ): device_capability_str = device_capability.as_version_str() - if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES: + if ( + device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES + and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str] + ): max_size = min( CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size], max_size, ) # device.index is a visible ordinal, not a logical local ID. - physical_device_id = current_platform.visible_device_id_to_physical_device_id( - device.index - ) - tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") - gather_list = [ - torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size) - ] - dist.all_gather(gather_list, tensor, group=self.group) - physical_device_ids = [t.item() for t in gather_list] - - # test nvlink first, this will filter out most of the cases - # where custom allreduce is not supported - # this checks hardware and driver support for NVLink - assert current_platform.is_cuda_alike() - fully_connected = current_platform.is_fully_connected(physical_device_ids) - if world_size > 2 and not fully_connected: + fully_connected = False + if same_node: + physical_device_id = ( + current_platform.visible_device_id_to_physical_device_id(device.index) + ) + tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") + gather_list = [ + torch.tensor([0], dtype=torch.int, device="cpu") + for _ in range(world_size) + ] + dist.all_gather(gather_list, tensor, group=self.group) + physical_device_ids = [t.item() for t in gather_list] + assert current_platform.is_cuda_alike() + fully_connected = current_platform.is_fully_connected(physical_device_ids) + if same_node and world_size > 2 and not fully_connected: logger.warning( "Custom allreduce is disabled because it's not supported on" " more than two PCIe-only GPUs. To silence this warning, " @@ -158,7 +187,11 @@ class CustomAllreduce: # this is expensive to compute at the first time # then we cache the result # On AMD GPU, p2p is always enabled between XGMI connected GPUs - if not current_platform.is_rocm() and not _can_p2p(rank, world_size): + if ( + same_node + and not current_platform.is_rocm() + and not _can_p2p(rank, world_size) + ): logger.warning( "Custom allreduce is disabled because your platform lacks " "GPU P2P capability or P2P test failed. To silence this " @@ -170,21 +203,40 @@ class CustomAllreduce: # Buffers memory are owned by this Python class and passed to C++. # Metadata composes of two parts: metadata for synchronization and a # temporary buffer for storing intermediate allreduce results. - self.meta_ptrs = self.create_shared_buffer( - ops.meta_size() + max_size, group=group, uncached=True - ) + if same_node: + self.meta_ptrs = self.create_shared_buffer( + ops.meta_size() + max_size, group=group, uncached=True + ) + else: + meta_ptr, _ = ops.allocate_shared_buffer_and_handle(ops.meta_size()) + self.meta_ptrs = [meta_ptr] * world_size # This is a pre-registered IPC buffer. In eager mode, input tensors - # are first copied into this buffer before allreduce is performed - self.buffer_ptrs = self.create_shared_buffer(max_size, group=group) - # This is a buffer for storing the tuples of pointers pointing to - # IPC buffers from all ranks. Each registered tuple has size of - # 8*world_size bytes where world_size is at most 8. Allocating 8MB - # is enough for 131072 such tuples. The largest model I've seen only - # needs less than 10000 of registered tuples. + # are first copied into this buffer before the operation is performed + legacy_buffer_size = max(max_size, max_all_gather_size, max_reduce_scatter_size) + if same_node: + self.buffer_ptrs = self.create_shared_buffer( + legacy_buffer_size, + group=group, + ) + else: + buffer_ptr, _ = ops.allocate_shared_buffer_and_handle(legacy_buffer_size) + self.buffer_ptrs = [buffer_ptr] * world_size + # This stores tuples of pointers to IPC buffers from all ranks. + # Each registered tuple contains at most 16 addresses. + # Allocating 8MB is enough for 65536 such tuples. The largest model uses + # fewer than 10000 registered tuples. self.rank_data = torch.empty( 8 * 1024 * 1024, dtype=torch.uint8, device=self.device ) self.max_size = max_size + self.max_all_gather_size = max_all_gather_size + if max_mnnvl_all_gather_size is None: + max_mnnvl_all_gather_size = self._DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES[ + world_size + ] + self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size + self.max_reduce_scatter_size = max_reduce_scatter_size + self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size self.rank = rank self.world_size = world_size self.fully_connected = fully_connected @@ -192,6 +244,72 @@ class CustomAllreduce: self.meta_ptrs, self.rank_data, rank, self.fully_connected ) ops.register_buffer(self._ptr, self.buffer_ptrs) + self._init_mnnvl_buffer( + max( + max_mnnvl_all_gather_size * world_size, + max_mnnvl_reduce_scatter_size, + ) + ) + if not same_node and not self.mnnvl_multicast_ptr: + logger.warning( + "Custom collectives are disabled because this multi-node " + "group does not support MNNVL multicast." + ) + self.close() + self.disabled = True + + def _init_mnnvl_buffer(self, stage_size: int) -> None: + if torch_symm_mem is None or not current_platform.is_cuda(): + return + try: + buffer_size = stage_size * 6 + buffer = torch_symm_mem.empty( + buffer_size, dtype=torch.uint8, device=self.device + ) + handle = torch_symm_mem.rendezvous(buffer, self.group.group_name) + if handle.multicast_ptr == 0: + return + peer_buffers = [ + handle.get_buffer( + peer, + (buffer_size,), + torch.uint8, + storage_offset=0, + ) + for peer in range(self.world_size) + ] + ptrs = [peer_buffer.data_ptr() for peer_buffer in peer_buffers] + lamport_ag_offset = 0 + lamport_rs_offset = stage_size * 3 + lamport_ag_ptrs = [ptr + lamport_ag_offset for ptr in ptrs] + lamport_rs_ptrs = [ptr + lamport_rs_offset for ptr in ptrs] + ops.register_buffer(self._ptr, lamport_ag_ptrs) + ops.register_buffer(self._ptr, lamport_rs_ptrs) + + buffer.view(torch.int32).fill_(-2147483648) + epochs = torch.zeros( + (2, 32), + dtype=torch.int32, + device=self.device, + ) + torch.accelerator.synchronize() + dist.barrier(group=self.group) + + self.mnnvl_buffer = buffer + self.mnnvl_handle = handle + self.mnnvl_peer_buffers = peer_buffers + self.mnnvl_multicast_ptr = handle.multicast_ptr + self.mnnvl_buffer_size = stage_size + self.mnnvl_lamport_ag_local_ptr = lamport_ag_ptrs[self.rank] + self.mnnvl_lamport_ag_multicast_ptr = ( + handle.multicast_ptr + lamport_ag_offset + ) + self.mnnvl_lamport_rs_local_ptr = lamport_rs_ptrs[self.rank] + self.mnnvl_lamport_epochs = epochs + self.mnnvl_lamport_ag_epoch_ptr = epochs[0].data_ptr() + self.mnnvl_lamport_rs_epoch_ptr = epochs[1].data_ptr() + except RuntimeError as error: + logger.debug("MNNVL AG/RS initialization failed: %s", error) @contextmanager def capture(self): @@ -228,7 +346,7 @@ class CustomAllreduce: ops.register_graph_buffers(self._ptr, handles, offsets) def should_custom_ar(self, inp: torch.Tensor): - if self.disabled: + if self.disabled or self.world_size > 8: return False inp_size = inp.numel() * inp.element_size() # custom allreduce requires input byte size to be multiples of 16 @@ -279,6 +397,111 @@ class CustomAllreduce: # latency) compared to the performance gain of using custom kernels return self.all_reduce(input, registered=False) + def should_custom_all_gather(self, inp: torch.Tensor) -> bool: + if self.disabled or not current_platform.is_cuda(): + return False + if self.world_size == 16 and not self.mnnvl_only: + return False + inp_size = inp.nbytes + if inp.dtype not in ( + torch.float32, + torch.float16, + torch.bfloat16, + ): + return False + max_size = ( + self.max_mnnvl_all_gather_size + if self.mnnvl_multicast_ptr + else self.max_all_gather_size + ) + return ( + 0 < inp_size <= max_size + and inp_size % 16 == 0 + and is_weak_contiguous(inp) + and (self.fully_connected or bool(self.mnnvl_multicast_ptr)) + ) + + def custom_all_gather(self, inp: torch.Tensor) -> torch.Tensor | None: + if not self.should_custom_all_gather(inp): + return None + out_shape = (inp.shape[0] * self.world_size,) + inp.shape[1:] + if self.mnnvl_multicast_ptr: + logger.info_once( + "Using the MNNVL Lamport all-gather kernel.", + scope="global", + ) + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + ops.mnnvl_lamport_all_gather( + self._ptr, + inp, + out, + self.mnnvl_lamport_ag_local_ptr, + self.mnnvl_lamport_ag_multicast_ptr, + self.mnnvl_lamport_ag_epoch_ptr, + self.mnnvl_buffer_size, + ) + else: + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + ops.custom_all_gather( + self._ptr, + inp, + out, + self.buffer_ptrs[self.rank], + self.max_all_gather_size, + ) + return out + + def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool: + if self.disabled or not current_platform.is_cuda(): + return False + if self.world_size == 16 and not self.mnnvl_only: + return False + inp_size = inp.nbytes + if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return False + if inp.shape[0] % self.world_size != 0: + return False + output_size = inp_size // self.world_size + max_size = ( + self.max_mnnvl_reduce_scatter_size + if self.mnnvl_multicast_ptr + else self.max_reduce_scatter_size + ) + return ( + 0 < inp_size <= max_size + and output_size % 16 == 0 + and is_weak_contiguous(inp) + and (self.fully_connected or bool(self.mnnvl_multicast_ptr)) + ) + + def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None: + if not self.should_custom_reduce_scatter(inp): + return None + out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:] + out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device) + if self.mnnvl_multicast_ptr: + logger.info_once( + "Using the MNNVL Lamport reduce-scatter kernel.", + scope="global", + ) + ops.mnnvl_lamport_reduce_scatter( + self._ptr, + inp, + out, + self.mnnvl_lamport_rs_local_ptr, + self.mnnvl_lamport_rs_epoch_ptr, + self.mnnvl_buffer_size, + ) + else: + ops.custom_reduce_scatter( + self._ptr, + inp, + out, + self.buffer_ptrs[self.rank], + self.max_reduce_scatter_size, + ) + return out + def close(self): if not self.disabled and self._ptr: if ops is not None: @@ -286,6 +509,10 @@ class CustomAllreduce: self._ptr = 0 self.free_shared_buffer(self.meta_ptrs, rank=self.rank) self.free_shared_buffer(self.buffer_ptrs, rank=self.rank) + self.mnnvl_peer_buffers = None + self.mnnvl_handle = None + self.mnnvl_buffer = None + self.mnnvl_lamport_epochs = None def __del__(self): self.close() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index fc9349a0be..fd1fa77e71 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -741,6 +741,7 @@ class EngineArgs: fail_on_environ_validation: bool = False gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None + kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None def __post_init__(self): # support `EngineArgs(compilation_config={...})` @@ -1637,6 +1638,13 @@ class EngineArgs: default=None, help="Select GDN prefill backend.", ) + parser.add_argument( + "--kda-prefill-backend", + dest="kda_prefill_backend", + choices=["auto", "triton", "flashkda"], + default=None, + help="Select KDA prefill backend.", + ) return parser @classmethod @@ -2431,6 +2439,8 @@ class EngineArgs: if self.gdn_prefill_backend is not None: self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend + if self.kda_prefill_backend is not None: + self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend config = VllmConfig( model_config=model_config, diff --git a/vllm/envs.py b/vllm/envs.py index 607b6ce25d..b3fbfccbc1 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -131,6 +131,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 + AITER_SITUV2_A8W4: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_USE_AITER_MHA: bool = True @@ -269,7 +270,6 @@ 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 @@ -1004,7 +1004,6 @@ environment_variables: dict[str, Callable[[], Any]] = { # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. - # - "identity": returns raw video bytes for model processor to handle. # # Custom backend implementations can be registered # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and @@ -1211,6 +1210,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_MOE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") ), + # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up- + # interleaved flydsl kernels instead of the default a16w4 separated path. + # Shared with the AITER runtime, which reads the same env var directly. + "AITER_SITUV2_A8W4": lambda: ( + os.getenv("AITER_SITUV2_A8W4", "0").lower() in ("true", "1") + ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass # for large batches @@ -1895,11 +1900,6 @@ 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/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 16fb961e56..4a265554b7 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -272,6 +272,7 @@ from vllm.v1.attention.backends.mla.prefill import ( ) from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, + get_num_attention_heads_from_layers, split_decodes_and_prefills, ) from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs @@ -375,6 +376,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): use_sparse: bool = False, indexer: object | None = None, topk_indices_buffer: torch.Tensor | None = None, + non_causal_multi_token_decode: bool = False, **extra_impl_args, ): super().__init__() @@ -391,6 +393,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): self.head_size = kv_lora_rank + qk_rope_head_dim self.layer_name = prefix self.indexer = indexer + self.non_causal_multi_token_decode = non_causal_multi_token_decode self.num_kv_heads = 1 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim @@ -1123,6 +1126,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): 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 _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): @@ -1450,6 +1454,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]): num_decode_tokens: int num_prefills: int + causal: bool = True + # The dimension of the attention heads head_dim: int | None = None @@ -1724,6 +1730,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # speculative decoding is enabled. query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY + # Whether this builder can flatten a non-causal query block into decode rows. + supports_non_causal_multi_token_decode: ClassVar[bool] = False + # The threshold for reordering the batch into decode and prefill requests. # If > 1, the batch will be reordered such that requests with # query length <= threshold are classified as decode requests. @@ -1824,8 +1833,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.vllm_config = vllm_config self.device = device self.use_pcp = parallel_config.prefill_context_parallel_size > 1 + self.non_causal_multi_token_decode = getattr( + kv_cache_spec, "non_causal_multi_token_decode", False + ) - self.num_heads = self.model_config.get_num_attention_heads(parallel_config) + # A draft cache group can have a different head count from the target. + self.num_heads = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or self.model_config.get_num_attention_heads(parallel_config) self.mla_dims = get_mla_dims(self.model_config) self.aot_schedule = current_platform.is_cuda() @@ -1951,14 +1966,42 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): seq_lens = common_attn_metadata.seq_lens dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills( - common_attn_metadata, - decode_threshold=self.reorder_batch_threshold, - require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), - treat_short_extends_as_decodes=not self.use_pcp, + non_causal_decode = common_attn_metadata.causal is False + if non_causal_decode: + if not ( + self.supports_non_causal_multi_token_decode + and self.non_causal_multi_token_decode + ): + raise ValueError( + "Non-causal multi-token MLA requires an explicitly supported " + "attention group." + ) + query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + num_active_reqs = int(torch.count_nonzero(query_lens > 0)) + uniform_active_queries = num_active_reqs > 0 and bool( + torch.all(query_lens[:num_active_reqs] == query_lens[0]) + ) + trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0)) + if not (uniform_active_queries and trailing_graph_padding): + raise ValueError( + "Non-causal MLA requires a uniform query block; got query " + f"lengths {query_lens.tolist()}." + ) + # Use exact GPU sequence lengths instead of the prefill path's CPU + # context-length upper bounds. + num_decodes = num_reqs + num_prefills = 0 + num_decode_tokens = num_tokens + num_prefill_tokens = 0 + else: + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), + treat_short_extends_as_decodes=not self.use_pcp, + ) ) - ) assert num_decodes + num_prefills == num_reqs assert num_decode_tokens + num_prefill_tokens == num_tokens @@ -2049,6 +2092,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, num_prefills=num_prefills, + causal=not non_causal_decode, prefill=prefill_metadata, decode=decode_metadata, ) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index f4319b99c8..27659d220d 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -505,7 +505,7 @@ def deepgemm_moe_permute( dtype=torch.int32, ) else: - aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32) + aq_scale_out = torch.zeros((M_sum, sf_k), device=device, dtype=torch.float32) # DeepGEMM uses negative values in m_indices (here expert_ids) to mark # completely invalid / padded blocks that should be skipped. We always diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index 90972fc7f0..ce30afd5be 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -204,6 +204,7 @@ def persistent_masked_m_silu_mul_quant( dtype=ys_dtype, device=y.device, ) + y_s.zero_() ceil_ue8m0 = quant_scale_fmt in [ DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, 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 6336154a9e..1125130c89 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -671,10 +671,10 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return not ( - moe_parallel_config.use_fi_nvl_two_sided_kernels - or moe_parallel_config.use_fi_nvl_one_sided_kernels - ) + # One-sided FI-NVL all2all pairs with MarlinExperts fine (the + # compressed-tensors MXFP4 path runs this exact combo); only the + # two-sided kernels are unsupported here. + return not moe_parallel_config.use_fi_nvl_two_sided_kernels @property def quant_type_id(self) -> int: @@ -1068,11 +1068,8 @@ class BatchedMarlinExperts(MarlinExpertsBase): return num_experts, max_num_tokens = hidden_states.shape[:2] - beta = activation_situ_beta + beta = 1.0 if activation_situ_beta is None else 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), diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 4f191334bb..bc770d0746 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -5,6 +5,7 @@ from functools import lru_cache import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -262,6 +263,8 @@ def rocm_aiter_fused_experts( elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu") activation_interleave = False + elif activation == MoEActivation.SITU: + activation_method = rocm_aiter_ops.get_aiter_activation_type("situ") else: raise ValueError(f"Unsupported activation: {activation}") @@ -356,10 +359,12 @@ def rocm_aiter_fused_experts( # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099 # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401 - hidden_pad = hidden_pad // 128 * 128 - intermediate_pad = ( - intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) - ) + # SITU's A16W4 FlyDSL kernel pads per gate/up half; pass through unrounded. + if activation != MoEActivation.SITU: + hidden_pad = hidden_pad // 128 * 128 + intermediate_pad = ( + intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) + ) # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs # for interleaved vs separated gate and up weights. @@ -370,7 +375,15 @@ def rocm_aiter_fused_experts( from aiter.ops.flydsl.moe_common import GateMode gate_mode = "" - if quant_config.use_mxfp4_w4a16: + if activation == MoEActivation.SITU: + # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) + # fp8 flydsl kernels; default a16w4 SiTU stays separated. + gate_mode = ( + GateMode.INTERLEAVE.value + if envs.AITER_SITUV2_A8W4 + else GateMode.SEPARATED.value + ) + elif quant_config.use_mxfp4_w4a16: gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: gate_mode = ( @@ -401,6 +414,8 @@ def rocm_aiter_fused_experts( bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + beta=moe_config.activation_situ_beta, + linear_beta=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 2027ed297c..315ec9f9c0 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 @@ -15,6 +15,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + activation_to_flashinfer_int, + has_flashinfer_situ_activation, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kMxfp4Static, @@ -133,19 +137,16 @@ class TrtLlmMxfp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: + if activation == MoEActivation.SITU: + return has_flashinfer_situ_activation() 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 + return activation_to_flashinfer_int(activation) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -315,6 +316,75 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula output = (M, self.hidden_dim_unpadded) return (workspace1, workspace2, output) + def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int: + """Max tokens per kernel call before the batched-GEMM grid overflows. + + The TRTLLM-Gen batched GEMM launches a static grid whose batch (Y) + dimension is ``getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts, + tileTokensDim)`` and must stay <= 65535. Solving that for num_tokens + with the smallest tile the kernel may pick (tileTokensDim=8, the runner + default) gives a bound that is safe regardless of the tactic selected. + Without it, large batches (e.g. Kimi-K3 top_k=16, EP16 profiling with + 131072 gathered tokens) overflow the grid and the GEMM launch fails. + """ + MAX_GRID_Y = 65535 + MIN_TILE_TOKENS_DIM = 8 + max_tokens = (MAX_GRID_Y - global_num_experts) * MIN_TILE_TOKENS_DIM // top_k + return max(1, min(300000, max_tokens)) + + def _invoke_kernel( + self, + output: torch.Tensor, + x_quant: torch.Tensor, + x_scale: torch.Tensor | None, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + local_num_experts: int, + local_expert_offset: int, + topk: int, + ) -> None: + from flashinfer import trtllm_fp4_block_scale_routed_moe + + packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + trtllm_fp4_block_scale_routed_moe( + topk_ids=packed_tensor, + routing_bias=None, + hidden_states=x_quant, + hidden_states_scale=x_scale, + gemm1_weights=w1, + gemm1_weights_scale=self.w1_scale, + gemm1_bias=self.w1_bias, + gemm1_alpha=self.gemm1_alpha, + gemm1_beta=self.gemm1_beta, + gemm1_clamp_limit=self.gemm1_clamp_limit, + gemm2_weights=w2, + gemm2_weights_scale=self.w2_scale, + gemm2_bias=self.w2_bias, + output1_scale_scalar=None, + output1_scale_gate_scalar=None, + output2_scale_scalar=None, + num_experts=global_num_experts, + top_k=topk, + n_group=None, + topk_group=None, + intermediate_size=self.intermediate_size_per_partition, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=None, + # Modular kernel receives pre-routed tokens, so routing is already + # done. Use Renormalize as a safe default the TRTLLM kernel supports. + 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), + ) + def apply( self, output: torch.Tensor, @@ -335,7 +405,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula ): topk = topk_ids.size(-1) local_num_experts = w1.size(0) - intermediate_size = self.intermediate_size_per_partition local_expert_offset = self.moe_config.ep_rank * local_num_experts if a1q_scale is not None: @@ -346,49 +415,27 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula x_quant = hidden_states x_scale = None - # Pack topk ids and weights into format expected by the kernel. - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) - assert self.w1_scale is not None assert self.w2_scale is not None - kwargs = { - "topk_ids": packed_tensor, - "routing_bias": None, - "hidden_states": x_quant, - "hidden_states_scale": x_scale, - "gemm1_weights": w1, - "gemm1_weights_scale": self.w1_scale, - "gemm1_bias": self.w1_bias, - "gemm1_alpha": self.gemm1_alpha, - "gemm1_beta": self.gemm1_beta, - "gemm1_clamp_limit": self.gemm1_clamp_limit, - "gemm2_weights": w2, - "gemm2_weights_scale": self.w2_scale, - "gemm2_bias": self.w2_bias, - "output1_scale_scalar": None, - "output1_scale_gate_scalar": None, - "output2_scale_scalar": None, - "num_experts": global_num_experts, - "top_k": topk, - "n_group": None, - "topk_group": None, - "intermediate_size": intermediate_size, - "local_expert_offset": local_expert_offset, - "local_num_experts": local_num_experts, - "routed_scaling_factor": None, - # Modular kernel receives pre-routed tokens, so routing - # is already done. Use Renormalize as a safe default that - # the TRTLLM C++ kernel supports. - "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), - } - from flashinfer import trtllm_fp4_block_scale_routed_moe - - trtllm_fp4_block_scale_routed_moe(**kwargs) + # Chunk tokens so the batched-GEMM grid stays within CUDA limits. + M = x_quant.size(0) + chunk_size = self._max_supported_tokens(topk, global_num_experts) + for start in range(0, M, chunk_size): + end = min(start + chunk_size, M) + self._invoke_kernel( + output[start:end], + x_quant[start:end], + None if x_scale is None else x_scale[start:end], + topk_ids[start:end], + topk_weights[start:end], + w1, + w2, + activation, + global_num_experts, + local_num_experts, + local_expert_offset, + topk, + ) return output 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 9bdff64281..a5555dc132 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 @@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.utils import ( ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, + has_flashinfer_situ_activation, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -225,15 +226,15 @@ class TrtLlmNvFp4ExpertsBase: @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the - private SituGLU (SITU) cubin.""" + """Supports SITU only when the installed FlashInfer exposes it.""" + if activation == MoEActivation.SITU: + return has_flashinfer_situ_activation() 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/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 07cff8fa22..0aed6cbf40 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -648,7 +648,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( else: hidden_size = round_up(hidden_size, 256) elif backend in TRTLLM_BACKENDS: - intermediate_size = round_up(intermediate_size, 256) + intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 256) elif backend in ( Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py index 257b5f04dc..d11beaf98f 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py @@ -102,7 +102,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): num_local_experts, self.max_num_tokens, hidden_dim ) - b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device) + b_a1_scale = torch.zeros(scale_shape, dtype=torch.float32, device=a1.device) else: assert quant_config.a1_scale is None b_a1_scale = None 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 index 5f88169221..9daac9fb53 100644 --- a/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py @@ -6,11 +6,6 @@ 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 @@ -82,25 +77,30 @@ class LatentMoERunner(MoERunner): ) self._k3_latent_moe_tail_op = op - def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor: + def _get_zero_residual( + self, + hidden_states: torch.Tensor, + max_token_num: int, + ) -> 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 - ): + if buf is None: buf = torch.zeros( - numel, dtype=hidden_states.dtype, device=hidden_states.device + max_token_num * hidden_states.shape[-1], + dtype=hidden_states.dtype, + device=hidden_states.device, ) self._zero_residual = buf - return buf[:numel].view_as(hidden_states) + + assert buf.dtype == hidden_states.dtype + assert buf.device == hidden_states.device + assert hidden_states.numel() <= buf.numel() + + return buf[: hidden_states.numel()].view_as(hidden_states) def _use_fused_path(self) -> bool: # The fused path merges the latent and shared reductions into one @@ -224,6 +224,12 @@ class LatentMoERunner(MoERunner): norm: RMSNorm, ) -> tuple[torch.Tensor, torch.Tensor]: """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.""" + from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + _AR_RESIDUAL_RMS_NORM, + _can_use_flashinfer, + flashinfer_trtllm_fused_allreduce_norm, + ) + if self.moe_config.tp_size == 1: return norm(hidden_states) @@ -238,7 +244,7 @@ class LatentMoERunner(MoERunner): # buffer and the normalized result into norm_out. flashinfer_trtllm_fused_allreduce_norm( allreduce_in=hidden_states, - residual=self._get_zero_residual(hidden_states), + residual=self._get_zero_residual(hidden_states, max_token_num), rms_gamma=norm.weight, rms_eps=norm.variance_epsilon, world_size=self.moe_config.tp_size, diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index f95aac54fd..bc49226c75 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -1,32 +1,33 @@ # 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.config import VllmConfig, get_current_vllm_config -from vllm.distributed import ( - divide, -) -from vllm.forward_context import ForwardContext, get_forward_context -from vllm.logger import init_logger +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.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention -from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader -from vllm.model_executor.utils import set_weight_attrs -from vllm.third_party.flash_linear_attention.ops.kda import ( - FusedRMSNormGated, - chunk_kda_with_fused_gate, - fused_kda_gate, - fused_recurrent_kda, +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.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.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from ...linear import ( ColumnParallelLinear, + MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) @@ -36,49 +37,118 @@ from ..mamba_utils import ( is_conv_state_dim_first, ) from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from ..ops.gather_initial_states import gather_initial_states -logger = init_logger(__name__) +# Empirical lower bound for the KDA gate to avoid numerical underflow. +_KDA_GATE_LOGBOUND_MIN = -5.0 -def kda_attention( - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, - g1: torch.Tensor, - beta: torch.Tensor, - core_attn_out: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self._forward( - q_proj_states=q_proj_states, - k_proj_states=k_proj_states, - v_proj_states=v_proj_states, - g1=g1, - beta=beta, - core_attn_out=core_attn_out, - ) +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 -def kda_attention_fake( - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, - g1: torch.Tensor, - beta: torch.Tensor, - core_attn_out: torch.Tensor, - layer_name: str, -) -> None: - return +def _make_fused_conv1d_weight_loader( + dims: list[int], + tp_size: int, + tp_rank: int, +) -> 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) + + return weight_loader -direct_register_custom_op( - op_name="kda_attention", - op_func=kda_attention, - mutates_args=["core_attn_out"], - fake_impl=kda_attention_fake, -) +class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): + """Merged projection with one output replicated across TP ranks. + + The replicated shard is represented as ``size * tp_size`` so the merged + parameter reserves ``size`` local rows on every rank. Loading that shard + from rank zero then gives every rank the complete checkpoint weight. + """ + + 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 @PluggableLayer.register("kimi_gated_delta_net_attention") @@ -96,7 +166,11 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): 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 + self.tp_size, + self.num_heads, + self.head_dim, + conv_kernel_size=self.conv_size, + num_spec=self.num_spec, ) def __init__( @@ -114,122 +188,142 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): assert self.num_heads % self.tp_size == 0 self.local_num_heads = divide(self.num_heads, self.tp_size) - projection_size = self.head_dim * self.num_heads + 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"] + self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False) - self.q_proj = ColumnParallelLinear( + if self.use_full_rank_gate: + # Keep f_a before the narrow beta shard, then pad each TP-local row + # to select the aligned BF16 GEMM path. The padding also avoids an + # Inductor correctness issue seen with the row-strided G view. + 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) + else: + in_proj_output_sizes = [self.projection_size] * 3 + [ + self.num_heads, + self.head_dim, + ] + self.in_proj_padding = 0 + self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( self.hidden_size, - projection_size, + in_proj_output_sizes, + replicated_shard_id=4, + tp_size=self.tp_size, bias=False, quant_config=self.quant_config, - prefix=f"{prefix}.q_proj", - ) - self.k_proj = ColumnParallelLinear( - self.hidden_size, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.k_proj", - ) - self.v_proj = ColumnParallelLinear( - self.hidden_size, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.v_proj", - ) - - self.f_a_proj = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.f_a_proj", + 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, - projection_size, + self.projection_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.f_b_proj", ) self.dt_bias = nn.Parameter( - torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32) + torch.empty(self.local_projection_size, dtype=torch.float32) ) set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - self.b_proj = ColumnParallelLinear( - self.hidden_size, - self.num_heads, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.b_proj", - ) - - self.q_conv1d = ColumnParallelLinear( + # 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=projection_size, + output_size=3 * self.projection_size, bias=False, params_dtype=torch.float32, - prefix=f"{prefix}.q_conv1d", + prefix=f"{prefix}.conv1d", ) - self.k_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.k_conv1d", + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": _make_fused_conv1d_weight_loader( + [self.projection_size] * 3, + self.tp_size, + self.tp_rank, + ) + }, ) - self.v_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.v_conv1d", - ) - # unsqueeze to fit conv1d weights shape into the linear weights shape. - # Can't do this in `weight_loader` since it already exists in - # `ColumnParallelLinear` and `set_weight_attrs` - # doesn't allow to override it - self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1) - self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) - self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) self.A_log = nn.Parameter( - torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) + torch.empty(self.local_num_heads, dtype=torch.float32) ) - set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)}) + set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)}) - self.g_a_proj = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.g_a_proj", + 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}." + ) + self.use_safe_gate = self.gate_lower_bound is not None + additional_config = vllm_config.additional_config + backend = ( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" ) - self.g_b_proj = ColumnParallelLinear( - self.head_dim, - projection_size, - bias=False, - quant_config=self.quant_config, - prefix=f"{prefix}.g_b_proj", + backend = "triton" if backend == "auto" else backend + assert backend == "triton", ( + "The shared Kimi GDN layer only supports the Triton KDA " + f"prefill backend, got {backend!r}." ) + if not self.use_full_rank_gate: + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + self.projection_size, + bias=False, + quant_config=self.quant_config, + prefix=f"{prefix}.g_b_proj", + ) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") self.o_proj = RowParallelLinear( - projection_size, + self.projection_size, self.hidden_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.o_proj", ) - compilation_config = get_current_vllm_config().compilation_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 + def rearrange_mixed_qkv( + self, mixed_qkv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + seq_len = mixed_qkv.shape[0] + qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim) + # Materialize all three row-strided inputs with one token-major to + # QKV-major permutation. Each unbound tensor is then contiguous. + qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1) + return qkv.unbind(0) + def forward( self, hidden_states: torch.Tensor, @@ -237,42 +331,57 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): output: torch.Tensor, ) -> None: num_tokens = hidden_states.size(0) - q = self.q_proj(hidden_states)[0] - k = self.k_proj(hidden_states)[0] - v = self.v_proj(hidden_states)[0] + projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0] + if self.use_full_rank_gate: + 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] + else: + mixed_qkv, beta, f_a = projected_qkvgfab.split( + [ + 3 * self.local_projection_size, + self.local_num_heads, + self.head_dim, + ], + dim=-1, + ) + g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] - beta = self.b_proj(hidden_states)[0].float().sigmoid() - g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] + 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) - g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) - core_attn_out = torch.zeros( + core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) - torch.ops.vllm.kda_attention( - q, - k, - v, - g1, - beta, - core_attn_out, - self.prefix, + + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, ) - core_attn_out = self.o_norm(core_attn_out, g2) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") output[:] = self.o_proj(core_attn_out)[0] + @eager_break_during_capture def _forward( self, - q_proj_states: torch.Tensor, - k_proj_states: torch.Tensor, - v_proj_states: torch.Tensor, + mixed_qkv: torch.Tensor, g1: torch.Tensor, + g2: torch.Tensor, beta: torch.Tensor, core_attn_out: torch.Tensor, ) -> None: @@ -280,165 +389,247 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention): attn_metadata_raw = forward_context.attn_metadata if attn_metadata_raw is None: - # # V1 profile run return + # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies + # under kimi_k3/{amd,nvidia}/ops so each can diverge independently. + if current_platform.is_rocm(): + from vllm.models.kimi_k3.amd.ops.third_party.kda import ( + chunk_kda_with_fused_gate, + fused_recurrent_kda, + fused_recurrent_kda_packed_decode, + ) + else: + 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, GDNAttentionMetadata) - has_initial_state = attn_metadata_narrowed.has_initial_state - non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc - non_spec_state_indices_tensor = ( - attn_metadata_narrowed.non_spec_state_indices_tensor - ) # noqa: E501 - num_actual_tokens = attn_metadata_narrowed.num_actual_tokens - constant_caches = self.kv_cache - - q_proj_states = q_proj_states[:num_actual_tokens] - k_proj_states = k_proj_states[:num_actual_tokens] - v_proj_states = v_proj_states[:num_actual_tokens] + 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_sequence_masks = m.spec_sequence_masks + 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 + mixed_qkv = mixed_qkv[:num_actual_tokens] g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] - (conv_state, recurrent_state) = constant_caches + constant_caches = self.kv_cache + + conv_state, recurrent_state = constant_caches # conv_state must be (..., dim, width-1) for the conv kernels. # DS layout stores it that way directly; SD layout needs a transpose. if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) - conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2) + 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 + ) - q_conv_weights = self.q_conv1d.weight.view( - self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2) - ) - k_conv_weights = self.k_conv1d.weight.view( - self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2) - ) - v_conv_weights = self.v_conv1d.weight.view( - self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2) - ) - if attn_metadata_narrowed.num_prefills > 0: - q_proj_states = q_proj_states.transpose(0, 1) - k_proj_states = k_proj_states.transpose(0, 1) - v_proj_states = v_proj_states.transpose(0, 1) - q = causal_conv1d_fn( - q_proj_states, - q_conv_weights, - self.q_conv1d.bias, - activation="silu", - conv_states=conv_state_q, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) - k = causal_conv1d_fn( - k_proj_states, - k_conv_weights, - self.k_conv1d.bias, - activation="silu", - conv_states=conv_state_k, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) - v = causal_conv1d_fn( - v_proj_states, - v_conv_weights, - self.v_conv1d.bias, - activation="silu", - conv_states=conv_state_v, - has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, - metadata=attn_metadata_narrowed, - ).transpose(0, 1) + # Split tokens into the multi-query spec-decode part and the remaining + # (prefill / plain decode) part. + if spec_sequence_masks is not None: + 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: + 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: - assert non_spec_state_indices_tensor is not None - decode_conv_indices = non_spec_state_indices_tensor[ - : attn_metadata_narrowed.num_actual_tokens - ] - q = causal_conv1d_update( - q_proj_states, - conv_state_q, - q_conv_weights, - self.q_conv1d.bias, - activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, - ) - k = causal_conv1d_update( - k_proj_states, - conv_state_k, - k_conv_weights, - self.k_conv1d.bias, - activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, - ) - v = causal_conv1d_update( - v_proj_states, - conv_state_v, - v_conv_weights, - self.v_conv1d.bias, - activation="silu", - conv_state_indices=decode_conv_indices, - validate_data=True, - ) + mixed_qkv_spec = g1_spec = beta_spec = None + mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta - q, k, v = map( - lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v) - ) + # ---------- spec-decode multi-query path ---------- + core_attn_out_spec = None + if spec_sequence_masks is not None: + 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) - if attn_metadata_narrowed.num_prefills > 0: - assert non_spec_state_indices_tensor is not None - assert has_initial_state is not None - zero_idx = non_spec_state_indices_tensor[~has_initial_state] - recurrent_state[zero_idx] = 0 - initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous() - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=g1, - beta=beta, + # Sibling beta and, for full-rank gates, output-gate views remain + # live, so write the convolution output separately. + spec_conv_out = torch.empty( + mixed_qkv_spec.shape, + dtype=mixed_qkv_spec.dtype, + device=mixed_qkv_spec.device, + ) + 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-only batches write directly into core_attn_out. + 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, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, - ) - # Init cache - recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state - else: - assert non_spec_query_start_loc is not None - g1 = fused_kda_gate( - rearrange(g1, "1 n h d -> n (h d)"), - self.A_log, - self.head_dim, - g_bias=self.dt_bias, - ).unsqueeze(0) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = fused_recurrent_kda( - q=q, - k=k, - v=v, - g=g1, - beta=beta, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, initial_state=recurrent_state, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc[ - : attn_metadata_narrowed.num_decodes + 1 - ], - ssm_state_indices=non_spec_state_indices_tensor, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, ) - core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ - 0, :num_actual_tokens - ] + + # ---------- non-spec path (prefill or plain decode) ---------- + 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 + ) + + # Packed prefill conv would require copying V solely to make + # it dense for KDA. Separate calls accept the strided inputs + # and produce dense Q/K/V without that extra traffic. + # TODO: Use packed conv once every KDA prefill backend accepts + # row-strided Q/K/V directly. + 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, + ) + ( + 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, + ) + # Init cache + recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + + else: + # pure-decode non-spec batch + assert non_spec_state_indices_tensor is not None + decode_conv_indices = non_spec_state_indices_tensor[ + : mixed_qkv_ns.size(0) + ] + # Sibling beta and, for full-rank gates, output-gate views + # remain live, so write the conv output separately. + packed_conv_out = torch.empty( + mixed_qkv_ns.shape, + dtype=mixed_qkv_ns.dtype, + device=mixed_qkv_ns.device, + ) + 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, + ) + + # ---------- merge spec and non-spec outputs ---------- + if core_attn_out_spec is not None and core_attn_out_non_spec is not None: + # Mixed batches require indexed placement in the original order. + merged = torch.empty( + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), + dtype=core_attn_out_spec.dtype, + device=core_attn_out_spec.device, + ) + merged.index_copy_(1, spec_token_indx, core_attn_out_spec) + merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens] + elif core_attn_out_non_spec is not None: + 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 + core_attn_out.copy_(self.o_norm(core_attn_out, g2)) diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index d974da5c27..e3361ca7dc 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -288,7 +288,7 @@ class MambaStateShapeCalculator: conv_dim = proj_size + 2 * proj_k_size conv_state_shape = cls._orient_conv_shape( - divide(conv_dim, tp_world_size), conv_kernel_size - 1 + divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec ) recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) return (conv_state_shape, recurrent_state_shape) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index ddeffb5645..8335c84966 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -8,11 +8,12 @@ import numpy as np import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID -@triton.jit() +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) def _causal_conv1d_fwd_kernel( # continuous batching # Pointers to matrices x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences @@ -33,7 +34,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, - num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + num_cache_lines, # added to support vLLM larger cache lines # Strides stride_x_dim: tl.constexpr, # stride to get to next feature-value, stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) @@ -58,6 +59,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, ): conv_states_ptr = initial_states_ptr conv_state_indices_ptr = cache_indices_ptr @@ -68,6 +70,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching KERNEL_WIDTH - 1 ) # can be passed via argument if it's not the same as this value + if launch_pdl: + tl.extra.cuda.gdc_wait() + # one program handles one chunk in a single sequence # rather than mixing sequences - to make updating initial_states across sequences efficiently @@ -79,6 +84,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) if idx_seq == pad_slot_id: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) @@ -136,6 +143,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching if HAS_NULL_BLOCK: # noqa if conv_states_input_coord == null_block_id: # not processing as this is a null block (padding) + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return conv_states_base = ( conv_states_ptr @@ -408,6 +417,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor w_col3 = tl.load(w_ptrs, mask_w, other=0.0) mask_x_1d = idx_feats < dim + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + for idx_token in range(segment_len): acc = acc_preload @@ -741,11 +754,12 @@ def causal_conv1d_fn( BLOCK_M=BLOCK_M, BLOCK_N=256, num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), ) return out.to(original_x_dtype) -@triton.jit() +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) def _causal_conv1d_update_kernel( # Pointers to matrices x_ptr, # (batch, dim, seqlen) @@ -763,7 +777,7 @@ def _causal_conv1d_update_kernel( dim: tl.constexpr, seqlen: tl.constexpr, state_len: tl.constexpr, - num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines + num_cache_lines, # added to support vLLM larger cache lines # Strides stride_x_seq: tl.constexpr, stride_x_dim: tl.constexpr, @@ -789,10 +803,16 @@ def _causal_conv1d_update_kernel( NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, ): + if launch_pdl: + tl.extra.cuda.gdc_wait() + # ruff: noqa: E501 idx_seq = tl.program_id(0) if idx_seq >= batch: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return # [BLOCK_N,] elements along the feature-dimension (channel) @@ -814,6 +834,8 @@ def _causal_conv1d_update_kernel( if HAS_NULL_BLOCK: # noqa if conv_states_input_coord == null_block_id: # not processing as this is not the actual sequence + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return if IS_VARLEN: @@ -831,6 +853,8 @@ def _causal_conv1d_update_kernel( o_offset = idx_seq * stride_o_seq if query_start_index == query_end_index: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() return if IS_SPEC_DECODING: @@ -969,6 +993,9 @@ def _causal_conv1d_update_kernel( mask_x_1d = idx_feats < dim # STEP 5: compute each token + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + for idx_token in tl.range(seqlen): acc = acc_preload @@ -1245,14 +1272,13 @@ def causal_conv1d_update( NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, BLOCK_N=256, + launch_pdl=current_platform.is_arch_support_pdl(), ) if unsqueeze: out = out.squeeze(-1) return out.to(original_x_dtype) -from vllm.platforms import current_platform # noqa: E402 - if current_platform.is_cpu(): from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_fn_cpu, diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index ac95646141..6c0e8e069f 100644 --- a/vllm/model_executor/layers/mla.py +++ b/vllm/model_executor/layers/mla.py @@ -28,6 +28,7 @@ class MLAModules: is_sparse: bool topk_indices_buffer: torch.Tensor | None indexer_rotary_emb: torch.nn.Module | None = None + g_proj: torch.nn.Module | None = None # --8<-- [start:multi_head_latent_attention] @@ -66,6 +67,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): quant_config: QuantizationConfig | None = None, prefix: str = "", skip_topk: bool = False, + non_causal_multi_token_decode: bool = False, allow_short_prefill_indexer_scoring_skip: bool = False, ) -> None: super().__init__() @@ -89,6 +91,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): self.indexer = mla_modules.indexer self.indexer_rope_emb = mla_modules.indexer_rotary_emb self.is_sparse = mla_modules.is_sparse + self.g_proj = mla_modules.g_proj # Whether to skip top-k token selection computation in this layer. # When True, the indexer will not be called, and the layer will reuse @@ -120,6 +123,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): use_sparse=self.is_sparse, indexer=self.indexer, topk_indices_buffer=mla_modules.topk_indices_buffer, + non_causal_multi_token_decode=non_causal_multi_token_decode, ) indexer_op = getattr(self.indexer, "indexer_op", None) if indexer_op is not None and hasattr( @@ -216,4 +220,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer): q_dcp_replicated=q_dcp_replicated, ) + if self.g_proj is not None: + attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid() + return self.o_proj(attn_out)[0] diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 5ef5fd40d5..76faec0e01 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -3,6 +3,7 @@ import torch +import vllm.envs as envs from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( @@ -17,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import modular_kernel as mk from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, + backend_to_kernel_cls, convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, @@ -68,8 +70,11 @@ class Mxfp4Config(QuantizationConfig): def get_config_filenames(cls) -> list[str]: return [] - # TODO (zyongye) This is only temporaty fallback. - # We should have `Mxfp4MoEMethod` after this migration is complete. + def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: + """MoE method for RoutedExperts. Subclasses override to pick a + checkpoint-specific kernel family.""" + return Mxfp4MoEMethod(moe) + def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": @@ -86,7 +91,7 @@ class Mxfp4Config(QuantizationConfig): ) return UnquantizedLinearMethod() elif isinstance(layer, RoutedExperts): - return GptOssMxfp4MoEMethod(layer.moe_config) + return self._make_moe_method(layer.moe_config) elif isinstance(layer, Attention): logger.debug_once( "MXFP4 attention layer is not implemented. " @@ -131,6 +136,9 @@ class GptOssMxfp4Config(Mxfp4Config): return None return "gpt_oss_mxfp4" + def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase: + return GptOssMxfp4MoEMethod(moe) + class GptOssMxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" @@ -475,13 +483,45 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): ) +def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool: + """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel. + + K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic + MXFP4 backend selector does not cover; route it to AITER on gfx950. + """ + from vllm.platforms import current_platform + + if not current_platform.is_rocm(): + return False + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.platforms.rocm import on_gfx950 + + return ( + rocm_aiter_ops.is_fused_moe_enabled() + and on_gfx950() + and moe.activation == MoEActivation.SITU + and moe.activation_situ_linear_beta is not None + and rocm_aiter_ops.get_aiter_activation_type("situ") is not None + ) + + class Mxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" def __init__(self, moe: FusedMoEConfig): super().__init__(moe) self.weight_dtype = "mxfp4" - self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) + self.is_k3_situ_aiter = _use_k3_situ_aiter(moe) + self.experts_cls: type[mk.FusedMoEExperts] | None + if self.is_k3_situ_aiter: + self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] + logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.") + else: + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend( + moe + ) self.max_capture_size = moe.max_capture_size @@ -523,6 +563,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): act_dtype=act_dtype, moe_parallel_config=moe_parallel_config, ) + if self.is_k3_situ_aiter: + # K3's AITER A16W4 kernel handles K3's native intermediate size + # (moe_intermediate 3072; e.g. 384/partition at TP8); the generic + # 256 round-up would inflate weights and OOM. + return hidden_size, intermediate_size_per_partition return mxfp4_round_up_hidden_size_and_intermediate_size( self.mxfp4_backend, hidden_size, intermediate_size_per_partition ) @@ -724,7 +769,58 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer=layer, ) + def _setup_kernel_k3_situ(self, layer: RoutedExperts) -> None: + # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all]) + # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in + # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here. + from aiter.utility.fp4_utils import e8m0_shuffle + + from vllm._aiter_ops import rocm_aiter_ops + + fp4_dtype = torch.float4_e2m1fn_x2 + e8m0_dtype = torch.float8_e8m0fnu + num_experts = layer.w13_weight.shape[0] + + # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) fp8 + # flydsl kernels, which need w13 weight+scale in interleave layout. + # Default a16w4 keeps the separated layout. + guinterleave = envs.AITER_SITUV2_A8W4 + w13 = rocm_aiter_ops.shuffle_weight_a16w4( + layer.w13_weight.data.view(fp4_dtype), 16, guinterleave + ) + w2 = rocm_aiter_ops.shuffle_weight_a16w4( + layer.w2_weight.data.view(fp4_dtype), 16, False + ) + w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype) + w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype) + w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( + w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave + ) + w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None and self.experts_cls is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + mxfp4_backend=self.mxfp4_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + def process_weights_after_loading(self, layer): + if self.is_k3_situ_aiter: + self._setup_kernel_k3_situ(layer) + return + w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale @@ -819,4 +915,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 632cca1fda..28c3bbae4c 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -18,20 +18,34 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int: return activation_to_flashinfer_type(activation).value +def has_flashinfer_situ_activation() -> bool: + try: + from flashinfer.fused_moe.core import ActivationType + except ImportError: + return False + return hasattr(ActivationType, "Situ") + + def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType": from flashinfer.fused_moe.core import ActivationType + if activation == MoEActivation.SITU: + situ = getattr(ActivationType, "Situ", None) + if situ is None: + raise ValueError("The installed FlashInfer does not support SITU") + return situ + # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively ACTIVATION_TO_FI_ACTIVATION = { MoEActivation.SILU_NO_MUL: ActivationType.Silu, MoEActivation.GELU_NO_MUL: ActivationType.Gelu, MoEActivation.SILU: ActivationType.Swiglu, # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args. + MoEActivation.SWIGLUOAI: ActivationType.Swiglu, MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, MoEActivation.GELU: ActivationType.Geglu, MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, - MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu, } return ACTIVATION_TO_FI_ACTIVATION[activation] diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index f7174ab8f1..cc8538e5ca 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -474,7 +474,11 @@ def prepare_moe_fp4_layer_for_marlin( group_size = 16 if is_nvfp4 else 32 - e = layer.moe_config.num_experts + # Use the per-rank (local) expert count: under expert parallelism the + # w13_weight/w2_weight tensors only hold this rank's experts (created with + # local_num_experts), whereas moe_config.num_experts is the global count. + # With no EP the two are equal, so the non-EP path is unchanged. + e = layer.moe_config.num_local_experts k = layer.moe_config.hidden_dim n = layer.moe_config.intermediate_size_per_partition diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index d076dce758..932bf92352 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -41,8 +41,8 @@ def _mxfp8_e4m3_quantize_torch( ) -> tuple[torch.Tensor, torch.Tensor]: """Naive MXFP8 quantization. For each block of 32 elements along the last dimension, compute a - shared e8m0 scale (the biased exponent of the block-wise amax) - and quantize each element to float8_e4m3fn. + shared e8m0 scale that fits the block-wise amax into the finite + float8_e4m3fn range, and quantize each element to float8_e4m3fn. Returns (quantized_values [same shape, fp8], scales uint8). Scale shape depends on is_sf_swizzled_layout: @@ -58,7 +58,8 @@ def _mxfp8_e4m3_quantize_torch( amax = x_blocked.abs().amax(dim=-1) amax = amax.clamp(min=torch.finfo(torch.float32).tiny) - scale_biased = torch.floor(torch.log2(amax)) + 127.0 + fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max + scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0 scale_biased = scale_biased.clamp(0, 254) scales_uint8 = scale_biased.to(torch.uint8) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 3367f4833e..7d97a3299d 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -16,10 +16,9 @@ import vllm.envs as envs from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import ( - Attention, - MLAAttention, MMEncoderAttention, ) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -122,12 +121,13 @@ def process_weights_after_loading( # the caching allocator, which starves the OS on UMA devices. release_device_memory_under_pressure(target_device) - # Initialize post-load attention weights for Attention, MLA, and MM encoder. - # NOTE: Happens after other modules so we can easily decompress weights. + # Initialize post-load attention weights for any attention layer and MM + # encoder. NOTE: Happens after other modules so we can easily decompress + # weights. for _, module in model.named_modules(): - if isinstance( - module, (Attention, MLAAttention, MMEncoderAttention) - ) and hasattr(module, "process_weights_after_loading"): + if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr( + module, "process_weights_after_loading" + ): # TODO(lucas): see if there is a way to unify the signatures # of process_weights_after_loading with device_loading_context(module, target_device): diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index ebcd26ac3f..05aa6a17dd 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -339,6 +339,39 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): ) +class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig): + """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface. + + Kimi-K3 ships its routed experts as compressed-tensors + ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which + lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel + selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config`` + (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4 + checkpoint is left untouched. Covers both the main model and the MTP draft. + + ``model_arch_config.quantization_config`` is a separate dict, snapshotted in + ``ModelConfig.__init__`` before this hook runs, and it is what + ``_verify_quantization`` reads when resolving the quant method. Patch it + alongside the hf configs so the rewrite lands before resolution; otherwise + the main model still resolves to compressed-tensors. + """ + + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + for cfg in ( + model_config.hf_config, + model_config.hf_text_config, + model_config.model_arch_config, + ): + quant_config = getattr(cfg, "quantization_config", None) + if ( + isinstance(quant_config, dict) + and quant_config.get("quant_method") == "compressed-tensors" + and quant_config.get("format") == "mxfp4-pack-quantized" + ): + quant_config["quant_method"] = "mxfp4" + + class GptOssForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -884,6 +917,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "JambaForSequenceClassification": JambaForSequenceClassificationConfig, "JinaForRanking": JinaForRankingConfig, "JinaVLForRanking": JinaVLForSequenceClassificationConfig, + "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig, + "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig, "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, "LlamaBidirectionalModel": LlamaBidirectionalConfig, "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 66068dcca6..2d1770e8d7 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -983,6 +983,7 @@ class DeepseekV2MLAAttention(nn.Module): topk_indices_buffer: torch.Tensor | None = None, input_size: int | None = None, reduce_results: bool = True, + non_causal_multi_token_decode: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -1178,6 +1179,7 @@ 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, + non_causal_multi_token_decode=non_causal_multi_token_decode, # 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, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 6556f6326e..668c958cef 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1407,7 +1407,7 @@ class EagleModelMixin: aux_hidden_states: list[torch.Tensor], layer_idx: int, hidden_states: torch.Tensor, - residual: torch.Tensor, + residual: torch.Tensor | None, ) -> list[torch.Tensor]: if layer_idx in self.aux_hidden_state_layers: value = hidden_states + residual if residual is not None else hidden_states diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index bcb8dd32de..1e5dd2cda3 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -34,6 +34,7 @@ from vllm.model_executor.models.vision import ( is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model, ) +from vllm.platforms import current_platform from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig logger = init_logger(__name__) @@ -154,9 +155,7 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module): def reset_parameters(self): nn.init.normal_(self.weight) - def forward( - self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] - ) -> torch.Tensor: + def get_pos_embeds(self, grid_thws: torch.Tensor | list[list[int]]) -> torch.Tensor: pos_embs = [] grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() for t, h, w in grid_thw_list: @@ -179,8 +178,12 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module): pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) - out = x + torch.cat(pos_embs) - return out + return torch.cat(pos_embs) + + def forward( + self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] + ) -> torch.Tensor: + return x + self.get_pos_embeds(grid_thws) class MoonVision3dPatchEmbed(nn.Module): @@ -195,6 +198,8 @@ class MoonVision3dPatchEmbed(nn.Module): pos_emb_width: int = 14, pos_emb_time: int = 4, pos_emb_type: str = "divided_fixed", + patch_embed_proj_bias: bool = True, + pos_emb_interpolation_mode: str = "bicubic", ): super().__init__() assert isinstance(patch_size, int | Sequence), ( @@ -208,7 +213,11 @@ class MoonVision3dPatchEmbed(nn.Module): self.patch_size = patch_size self.proj = nn.Conv2d( - in_dim, out_dim, kernel_size=patch_size, stride=patch_size + in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias, ) if pos_emb_type == "divided_fixed": @@ -217,17 +226,37 @@ class MoonVision3dPatchEmbed(nn.Module): width=pos_emb_width, num_frames=pos_emb_time, dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode, ) else: raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") def forward( - self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]] + self, + x: torch.Tensor, + grid_thws: torch.Tensor | list[list[int]] | None, + *, + pos_embeds: torch.Tensor | None = None, ) -> torch.Tensor: - x = self.proj(x).view(x.size(0), -1) - # apply positional embedding - x = self.pos_emb(x, grid_thws) - return x + x = self._proj(x).view(x.size(0), -1) + if pos_embeds is not None: + return x + pos_embeds + assert grid_thws is not None + return self.pos_emb(x, grid_thws) + + def _proj(self, x: torch.Tensor) -> torch.Tensor: + # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton. + if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16): + from aiter.ops.triton.conv.conv2d import conv2d + + return conv2d( + x, + self.proj.weight, + self.proj.bias, + stride=self.patch_size, + layout="nchw", + ) + return self.proj(x) class Rope2DPosEmbRepeated(nn.Module): @@ -342,6 +371,14 @@ class MLP2(nn.Module): return x +def _make_vision_norm(norm_type: str, hidden_dim: int) -> nn.Module: + if norm_type == "layernorm": + return nn.LayerNorm(hidden_dim) + if norm_type == "rmsnorm": + return nn.RMSNorm(hidden_dim) + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + class MoonViTEncoderLayer(nn.Module): """Single encoder layer for MoonViT with TP/DP support.""" @@ -355,23 +392,33 @@ class MoonViTEncoderLayer(nn.Module): *, activation=F.gelu, attn_bias: bool = False, + qkv_hidden_size: int | None = None, + norm_type: str = "layernorm", + mlp_type: str = "mlp2", + linear_bias: bool = True, ): super().__init__() self.use_data_parallel = is_vit_use_data_parallel(num_heads) self.num_heads = num_heads self.hidden_dim = hidden_dim - self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads + self.qkv_hidden_size = ( + hidden_dim if qkv_hidden_size is None else qkv_hidden_size + ) + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads self.tp_size = ( 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() ) self.num_attention_heads_per_partition = divide(num_heads, self.tp_size) - self.norm0 = nn.LayerNorm(hidden_dim) - self.norm1 = nn.LayerNorm(hidden_dim) + self.norm0 = _make_vision_norm(norm_type, hidden_dim) + self.norm1 = _make_vision_norm(norm_type, hidden_dim) + if mlp_type != "mlp2": + raise NotImplementedError(f"Not support mlp_type: {mlp_type}") self.mlp = MLP2( [hidden_dim, mlp_dim, hidden_dim], activation, + bias=linear_bias, quant_config=quant_config, prefix=f"{prefix}.mlp", use_data_parallel=self.use_data_parallel, @@ -387,7 +434,7 @@ class MoonViTEncoderLayer(nn.Module): disable_tp=self.use_data_parallel, ) self.wo = RowParallelLinear( - hidden_dim, + self.qkv_hidden_size, hidden_dim, bias=attn_bias, quant_config=quant_config, @@ -493,8 +540,9 @@ class MoonViT3dEncoder(nn.Module): f'video_attn_type must be "spatial_temporal", got {video_attn_type}' ) self.video_attn_type = video_attn_type + qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"] self.rope_2d = Rope2DPosEmbRepeated( - block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 + qkv_hidden_size // block_cfg["num_heads"], 512, 512 ) self.blocks = nn.ModuleList( [ @@ -506,13 +554,17 @@ class MoonViT3dEncoder(nn.Module): for layer_idx in range(num_layers) ] ) - self.final_layernorm = nn.LayerNorm(hidden_dim) + self.final_layernorm = _make_vision_norm( + block_cfg.get("norm_type", "layernorm"), hidden_dim + ) def prepare_encoder_metadata( self, grid_thw_list: list[list[int]], *, device: torch.device, + max_batch_size: int | None = None, + max_seqlen_override: int | None = None, ) -> dict[str, torch.Tensor | None]: metadata: dict[str, torch.Tensor | None] = {} metadata["rope_freqs_cis"] = self.rope_2d.get_freqs_cis( @@ -524,15 +576,30 @@ class MoonViT3dEncoder(nn.Module): cu_seqlens = np.concatenate( [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)] ) + if max_batch_size is not None: + num_seqs = len(cu_seqlens) - 1 + if num_seqs < max_batch_size: + cu_seqlens = np.concatenate( + [ + cu_seqlens, + np.full( + max_batch_size - num_seqs, + cu_seqlens[-1], + dtype=np.int32, + ), + ] + ) attn_backend = self.blocks[0].attn.attn_backend metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens( attn_backend, cu_seqlens, device ) - metadata["max_seqlen"] = torch.tensor( - MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens), - dtype=torch.int32, + max_seqlen = ( + max_seqlen_override + if max_seqlen_override is not None + else MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens) ) + metadata["max_seqlen"] = torch.tensor(max_seqlen, dtype=torch.int32) metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens( attn_backend, cu_seqlens, @@ -545,11 +612,12 @@ class MoonViT3dEncoder(nn.Module): def forward( self, hidden_states: torch.Tensor, - grid_thws: torch.Tensor | list[list[int]], + grid_thws: torch.Tensor | list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: if encoder_metadata is None: + assert grid_thws is not None grid_thw_list = ( grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() ) @@ -604,6 +672,34 @@ def tpool_patch_merger( return outputs +def build_image_merge_gather_idx( + grid_thws: list[list[int]] | list[tuple[int, int, int]], + merge_kernel_size: tuple[int, int], +) -> np.ndarray: + """Build packed spatial-merge indices for image-only CUDA graphs.""" + kh, kw = merge_kernel_size + parts: list[np.ndarray] = [] + offset = 0 + for t, h, w in grid_thws: + if t != 1: + raise ValueError("Image encoder CUDA graphs require grid T == 1") + idx = np.arange(h * w, dtype=np.int64).reshape(h, w) + idx = idx.reshape(h // kh, kh, w // kw, kw) + parts.append(idx.transpose(0, 2, 1, 3).reshape(-1, kh * kw) + offset) + offset += h * w + if not parts: + return np.empty((0, kh * kw), dtype=np.int64) + return np.concatenate(parts) + + +def tpool_patch_merger_packed( + x: torch.Tensor, + merge_gather_idx: torch.Tensor, +) -> torch.Tensor: + """Apply the image-only spatial merge using precomputed tensor indices.""" + return x[merge_gather_idx] + + class MoonViT3dPretrainedModel(nn.Module): """Main vision tower model. @@ -630,6 +726,10 @@ class MoonViT3dPretrainedModel(nn.Module): pos_emb_width=config.init_pos_emb_width, pos_emb_time=config.init_pos_emb_time, pos_emb_type=config.pos_emb_type, + patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), + pos_emb_interpolation_mode=getattr( + config, "pos_emb_interpolation_mode", "bicubic" + ), ) self.encoder = MoonViT3dEncoder( @@ -638,9 +738,15 @@ class MoonViT3dPretrainedModel(nn.Module): block_cfg={ "num_heads": config.num_attention_heads, "hidden_dim": config.hidden_size, + "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), "mlp_dim": config.intermediate_size, - "activation": get_act_fn("gelu_pytorch_tanh"), - "attn_bias": True, + "activation": get_act_fn( + getattr(config, "activation_func", "gelu_pytorch_tanh") + ), + "attn_bias": getattr(config, "attn_bias", True), + "norm_type": getattr(config, "norm_type", "layernorm"), + "mlp_type": getattr(config, "mlp_type", "mlp2"), + "linear_bias": getattr(config, "linear_bias", True), }, video_attn_type=config.video_attn_type, quant_config=quant_config, @@ -650,7 +756,7 @@ class MoonViT3dPretrainedModel(nn.Module): def forward( self, pixel_values: torch.Tensor, - grid_thws: torch.Tensor | list[list[int]], + grid_thws: torch.Tensor | list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: @@ -662,6 +768,22 @@ class MoonViT3dPretrainedModel(nn.Module): Returns: torch.Tensor: The output tokens. """ + if encoder_metadata is not None and "pos_embeds" in encoder_metadata: + hidden_states = self.patch_embed( + pixel_values, + None, + pos_embeds=encoder_metadata["pos_embeds"], + ) + hidden_states = self.encoder( + hidden_states, + None, + encoder_metadata=encoder_metadata, + ) + merge_gather_idx = encoder_metadata["merge_gather_idx"] + assert merge_gather_idx is not None + return tpool_patch_merger_packed(hidden_states, merge_gather_idx) + + assert grid_thws is not None grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() if encoder_metadata is None: encoder_metadata = self.encoder.prepare_encoder_metadata( @@ -685,13 +807,40 @@ class MoonViT3dPretrainedModel(nn.Module): return hidden_states + def prepare_encoder_cudagraph_metadata( + self, + grid_thw_list: list[list[int]], + *, + max_batch_size: int, + max_seqlen_override: int | None = None, + device: torch.device, + ) -> dict[str, torch.Tensor | None]: + """Precompute fixed-buffer metadata for image encoder CUDA graphs.""" + grid_thw_list = [list(map(int, grid)) for grid in grid_thw_list] + metadata = self.encoder.prepare_encoder_metadata( + grid_thw_list, + device=device, + max_batch_size=max_batch_size, + max_seqlen_override=max_seqlen_override, + ) + metadata["pos_embeds"] = self.patch_embed.pos_emb.get_pos_embeds( + grid_thw_list + ).to(device=device) + merge_gather_idx = build_image_merge_gather_idx( + grid_thw_list, self.merge_kernel_size + ) + metadata["merge_gather_idx"] = torch.from_numpy(merge_gather_idx).to( + device=device, non_blocking=True + ) + return metadata + @torch.inference_mode() def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]): """Apply MM projector to vision tower outputs.""" num_embedding_list = [x.shape[0] for x in vt_output] batched = torch.cat(vt_output, dim=0) - projector_dtype = mm_projector.pre_norm.weight.dtype + projector_dtype = next(mm_projector.parameters()).dtype if batched.dtype != projector_dtype: batched = batched.to(projector_dtype) proj_out = mm_projector(batched) @@ -747,11 +896,34 @@ class KimiK25MultiModalProjector(nn.Module): ): super().__init__() self.use_data_parallel = use_data_parallel + self.mm_projector_type = getattr(config, "mm_projector_type", "patchmerger") # Hidden size after patch merging merge_h, merge_w = config.merge_kernel_size self.hidden_size = config.hidden_size * merge_h * merge_w + if self.mm_projector_type == "patchmergerv2": + self.linear_1 = ReplicatedLinear( + self.hidden_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.linear_1", + ) + self.linear_2 = ReplicatedLinear( + self.hidden_size, + getattr(config, "text_hidden_size", config.mm_hidden_size), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.linear_2", + ) + self.post_norm = torch.nn.RMSNorm( + getattr(config, "text_hidden_size", config.mm_hidden_size), + eps=config.projector_ln_eps, + ) + self.act = GELUActivation() + return + self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-5) self.linear_1 = ReplicatedLinear( self.hidden_size, @@ -770,6 +942,13 @@ class KimiK25MultiModalProjector(nn.Module): self.act = GELUActivation() def forward(self, image_features: torch.Tensor) -> torch.Tensor: + if self.mm_projector_type == "patchmergerv2": + hidden_states = image_features.view(image_features.shape[0], -1) + hidden_states, _ = self.linear_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return self.post_norm(hidden_states) + hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size) hidden_states, _ = self.linear_1(hidden_states) hidden_states = self.act(hidden_states) diff --git a/vllm/model_executor/models/kimi_linear.py b/vllm/model_executor/models/kimi_linear.py deleted file mode 100644 index 057d4d01cb..0000000000 --- a/vllm/model_executor/models/kimi_linear.py +++ /dev/null @@ -1,646 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from collections.abc import Iterable - -import torch -from torch import nn - -from vllm.compilation.decorators import support_torch_compile -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 -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - fused_moe_make_expert_params_mapping, -) -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.sequence import IntermediateTensors -from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig - -from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - get_spec_layer_idx_from_weight_name, - is_pp_missing_parameter, - make_layers, - maybe_prefix, -) - -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 = "", - ) -> 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": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - 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 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 - intermediate_size = config.intermediate_size - moe_intermediate_size = config.moe_intermediate_size - num_experts = config.num_experts - moe_renormalize = config.moe_renormalize - 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 - - if config.hidden_act != "silu": - raise ValueError( - f"Unsupported activation: {config.hidden_act}. " - "Only silu is supported for now." - ) - - # Gate always runs at half / full precision for now. - self.gate = ReplicatedLinear( - hidden_size, - num_experts, - bias=False, - quant_config=None, - prefix=f"{prefix}.gate", - ) - - self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts)) - - if self.num_shared_experts is not None: - intermediate_size = moe_intermediate_size * self.num_shared_experts - self.shared_experts = KimiMLP( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=False, - prefix=f"{prefix}.shared_experts", - ) - else: - self.shared_experts = None - - self.experts = FusedMoE( - shared_experts=self.shared_experts, - num_experts=num_experts, - top_k=config.num_experts_per_token, - hidden_size=hidden_size, - intermediate_size=moe_intermediate_size, - 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, - ) - - 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 self.q_lora_rank is None - assert num_heads % tp_size == 0 - 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.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", - ) - - mla_modules = MLAModules( - 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=None, - kv_a_proj_with_mqa=self.kv_a_proj_with_mqa, - q_a_layernorm=None, - q_b_proj=None, - q_proj=self.q_proj, - indexer=None, - is_sparse=False, - topk_indices_buffer=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.is_moe = config.is_moe - layer_idx = int(prefix.rsplit(".", 1)[1]) - 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: - 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=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, - use_nope=config.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", - ) - 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", - ) - 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, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - # 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) - - attn_output = torch.empty_like(hidden_states) - self.self_attn( - hidden_states=hidden_states, - positions=positions, - output=attn_output, - ) - hidden_states = attn_output - - # Fully Connected - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -@support_torch_compile -class KimiLinearModel(nn.Module): - 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) - else: - self.norm = 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 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: - 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"] - - for _, layer in enumerate(self.layers[self.start_layer : self.end_layer]): - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 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()) - loaded_params: set[str] = set() - for args in weights: - name, loaded_weight = args[:2] - kwargs = args[2] if len(args) > 2 else {} - if "rotary_emb.inv_freq" in name: - continue - - 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 = name.replace(weight_name, param_name) - # 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 idx, (param_name, weight_name, expert_id, shard_id) in enumerate( - expert_params_mapping - ): - if weight_name not in name: - continue - name = name.replace(weight_name, 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=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. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - 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 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, ...], tuple[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: - 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/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index ca093e5968..db3464a446 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -137,7 +137,10 @@ _TEXT_GENERATION_MODELS = { "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), "JambaForCausalLM": ("jamba", "JambaForCausalLM"), - "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), + "KimiLinearForCausalLM": ( + "vllm.models.kimi_k3", + "KimiLinearForCausalLM", + ), "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"), "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"), "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"), @@ -460,6 +463,10 @@ _MULTIMODAL_MODELS = { ), "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), + "KimiK3ForConditionalGeneration": ( + "vllm.models.kimi_k3", + "KimiK3ForConditionalGeneration", + ), "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), "MossTranscribeDiarizeForConditionalGeneration": ( "moss_transcribe_diarize", @@ -609,6 +616,10 @@ _SPECULATIVE_DECODING_MODELS = { "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), + "K3DSparkModel": ( + "vllm.models.kimi_k3.nvidia.dspark_mla", + "K3DSparkForCausalLM", + ), "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"), "Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"), "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"), @@ -649,6 +660,7 @@ _SPECULATIVE_DECODING_MODELS = { "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"), + "KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"), # Temporarily disabled. # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1. # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"), diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 4b05cad1b1..cccc73a28a 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -234,7 +234,7 @@ def _deepgemm_fp8_gemm_nt_warmup( device = w.device a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.empty( + a1q_scales = torch.zeros( (max_tokens, k // block_m), device=device, dtype=torch.float32 ) out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16) @@ -336,7 +336,7 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( def _warmup(w: torch.Tensor, w_scale: torch.Tensor): _, n, k = w.size() a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn) - a1q_scales = torch.empty( + a1q_scales = torch.zeros( (MAX_M, k // block_m), device=device, dtype=torch.float32 ) out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index afe47f39ce..537a2311fb 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -28,6 +28,9 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( deepseek_v4_sparse_mla_attention_warmup, flashinfer_sparse_mla_decode_autotune_warmup, ) +from vllm.model_executor.warmup.kimi_k3_triton_warmup import ( + kimi_k3_triton_warmup, +) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( sparse_mla_triton_warmup, @@ -123,6 +126,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): # Run next so input-prep kernels JIT against pristine runner state. if worker.vllm_config.kernel_config.enable_jit_warmup: + kimi_k3_triton_warmup(worker) fa4_cutedsl_warmup(worker) sparse_mla_triton_warmup(worker) diff --git a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py new file mode 100644 index 0000000000..4c4a0327c7 --- /dev/null +++ b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up Kimi-K3 Triton kernels.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None: + from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention + + compilation_config = getattr( + worker.model_runner, + "compilation_config", + None, + ) + static_context = getattr(compilation_config, "static_forward_context", None) + if not isinstance(static_context, dict): + return None + return next( + ( + layer + for layer in static_context.values() + if isinstance(layer, KimiK3DeltaAttention) + ), + None, + ) + + +def _warm_attn_res(worker: Worker) -> None: + from vllm.models.kimi_k3.nvidia.ops.attn_res import ( + attn_res, + get_attn_res_triton_warmup_profiles, + ) + + config = worker.model_config.hf_text_config + block_size = getattr(config, "attn_res_block_size", None) + if block_size is None: + return + + hidden_size = int(config.hidden_size) + max_blocks = (int(config.num_hidden_layers) + block_size - 1) // block_size + if max_blocks < 2: + return + + dtype = worker.model_config.dtype + device = torch.device("cuda") + eps = float(config.rms_norm_eps) + prefix = torch.zeros((1, hidden_size), dtype=dtype, device=device) + delta = torch.zeros_like(prefix) + blocks = torch.zeros( + (1, max_blocks, hidden_size), + dtype=dtype, + device=device, + ) + norm_weight = torch.zeros(hidden_size, dtype=dtype, device=device) + qk_weight = torch.zeros_like(norm_weight) + output_norm_weight = torch.zeros_like(norm_weight) + + for ( + num_blocks, + has_delta, + block_write_idx, + apply_output_norm, + ) in get_attn_res_triton_warmup_profiles(max_blocks): + attn_res( + prefix, + delta if has_delta else None, + blocks, + norm_weight, + qk_weight, + output_norm_weight if apply_output_norm else None, + num_blocks=num_blocks, + block_write_idx=block_write_idx, + eps=eps, + output_norm_eps=eps if apply_output_norm else 0.0, + ) + + +def _warm_recurrent_kda( + layer: KimiK3DeltaAttention, + input_dtype: torch.dtype, +) -> None: + from vllm.models.kimi_k3.nvidia.ops.third_party.kda.fused_recurrent import ( + fused_recurrent_kda, + get_fused_recurrent_kda_fwd_warmup_profiles, + ) + + num_speculative_tokens = int(layer.num_spec) + # fused_recurrent_kda_fwd_kernel is only used by speculative decode. + if num_speculative_tokens <= 0: + return + + kv_cache = layer.kv_cache + if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2: + return + state = kv_cache[1] + if not isinstance(state, torch.Tensor) or not state.numel(): + return + + logger.info("Warming up Kimi-K3 speculative KDA kernels.") + h = int(layer.local_num_heads) + d = int(layer.head_dim) + tokens_per_sequence = num_speculative_tokens + 1 + for num_sequences in get_fused_recurrent_kda_fwd_warmup_profiles(h): + num_tokens = num_sequences * tokens_per_sequence + packed_qkv = torch.empty( + (num_tokens, 3 * h * d), + dtype=input_dtype, + device=state.device, + ) + q, k, v = ( + tensor.view(1, num_tokens, h, d) + for tensor in packed_qkv.split(h * d, dim=-1) + ) + fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_g=torch.empty( + (1, num_tokens, h, d), + dtype=input_dtype, + device=state.device, + ), + raw_beta=torch.empty( + (1, num_tokens, h), + dtype=input_dtype, + device=state.device, + ), + A_log=layer.A_log, + dt_bias=layer.dt_bias, + lower_bound=layer.gate_lower_bound, + initial_state=state[:1], + cu_seqlens=torch.arange( + 0, + num_tokens + 1, + tokens_per_sequence, + dtype=torch.int32, + device=state.device, + ), + ssm_state_indices=torch.zeros( + (num_sequences, tokens_per_sequence), + dtype=torch.int32, + device=state.device, + ), + num_accepted_tokens=torch.ones( + num_sequences, + dtype=torch.int32, + device=state.device, + ), + out=torch.empty( + (1, num_tokens, h, d), + dtype=input_dtype, + device=state.device, + ), + ) + + +@torch.inference_mode() +def kimi_k3_triton_warmup(worker: Worker) -> None: + """Warm Kimi-K3 Triton kernels reachable by this server.""" + if not current_platform.is_cuda(): + return + + layer = _get_kda_layer(worker) + if layer is None: + return + + _warm_attn_res(worker) + _warm_recurrent_kda(layer, worker.model_config.dtype) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 346b196b3c..5d89337ba7 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -22,9 +22,9 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v4.common.ops import ( fused_indexer_q_rope_quant, - fused_q_kv_rmsnorm, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE diff --git a/vllm/models/deepseek_v4/common/ops/__init__.py b/vllm/models/deepseek_v4/common/ops/__init__.py index ff6ee22996..b7899f01de 100644 --- a/vllm/models/deepseek_v4/common/ops/__init__.py +++ b/vllm/models/deepseek_v4/common/ops/__init__.py @@ -11,7 +11,6 @@ from .cache_utils import ( from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm -from .fused_qk_rmsnorm import fused_q_kv_rmsnorm from .save_partial_states import save_partial_states __all__ = [ @@ -23,7 +22,6 @@ __all__ = [ "fused_indexer_q_rope_quant", "fused_inv_rope_fp8_quant", "fused_mtp_input_rmsnorm", - "fused_q_kv_rmsnorm", "mtp_shared_head_rmsnorm", "quantize_and_insert_k_cache", "save_partial_states", diff --git a/vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py deleted file mode 100644 index 0dd348a46e..0000000000 --- a/vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch - -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, -): - # 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 - - # 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, - ) - return qr_out, kv_out diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index f9692d5d47..33a4012443 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -25,26 +25,6 @@ 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, @@ -148,7 +128,7 @@ class K3DSparkModel(nn.Module): ] ) self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) - self.markov_head = ReplicatedDSparkMarkovHead( + self.markov_head = DSparkMarkovHead( self.config.vocab_size, self.config.draft_vocab_size, self.config.markov_rank, diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index ec2dc651fa..ff7d06d3ff 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -121,6 +121,13 @@ from ..common.mm_preprocess import ( logger = init_logger(__name__) +# Token-count cutoff for overlapping the MoE router gate with the routed-expert +# down projection on a separate CUDA stream (latent MoE). At or below this many +# tokens the launch-bound decode path benefits from multi-stream overlap; above +# it the GEMMs saturate the device and the cross-stream sync is pure overhead, +# so it falls back to sequential. +_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256 + class KimiMLP(nn.Module): def __init__( @@ -513,7 +520,7 @@ class KimiMoE(nn.Module): ) # 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). + # _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: @@ -546,7 +553,13 @@ class KimiMoE(nn.Module): activation_linear_beta=activation_situ_linear_beta, ) else: - enable_tail_fusion = envs.VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION + # The tail-fusion kernels are tcgen05-based, so they require an + # SM100 NVIDIA device; the runner falls back to the default latent + # MoE path everywhere else. + enable_tail_fusion = ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) self.experts = FusedMoE( shared_experts=self.shared_experts, num_experts=num_experts, @@ -645,7 +658,7 @@ class KimiMoE(nn.Module): 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 + if num_tokens <= _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD else None, ) ) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 07dce7d138..f0b93da292 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -880,6 +880,12 @@ class Platform: ), cache_config.block_size, ) + if model_config.use_mla: + # TRTLLM/FlashInfer MLA decode kernels require the physical + # number of kernel blocks to be aligned to 128 / kernel_block_size. + # For hybrid MLA/Mamba models, make the manager block size a + # multiple of 128 so split kernel blocks keep that invariant. + kernel_block_alignment_size = max(kernel_block_alignment_size, 128) if cache_config.mamba_cache_mode == "all": # With prefix caching, align to mamba chunk size for kernel perf diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index 920efa4441..c8004cb571 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -12,7 +12,7 @@ import torch from vllm.triton_utils import tl, triton -from .op import exp +from .op import exp, log @triton.heuristics( diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index b1aaa294c7..89da81ea91 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", + k3_dspark="K3DSparkConfig", flex_olmo="FlexOlmoConfig", fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", @@ -100,6 +101,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", kimi_k25="KimiK25Config", + kimi_k3="KimiK3Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) mlp_speculator="MLPSpeculatorConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 4bb7674ddb..6685ee4f78 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", + "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark", "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr", "EAGLEConfig": "vllm.transformers_utils.configs.eagle", "FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid", @@ -72,6 +73,8 @@ _CLASS_TO_MODULE: dict[str, str] = { "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", + "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", + "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", "NemotronConfig": "vllm.transformers_utils.configs.nemotron", "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h", "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid", @@ -124,6 +127,7 @@ __all__ = [ "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", + "K3DSparkConfig", "DotsOCRConfig", "EAGLEConfig", "FlexOlmoConfig", @@ -156,6 +160,8 @@ __all__ = [ "KimiLinearConfig", "KimiVLConfig", "KimiK25Config", + "KimiK3Config", + "KimiK3VisionConfig", "NemotronConfig", "NemotronHConfig", "OlmoHybridConfig", diff --git a/vllm/transformers_utils/configs/k3_dspark.py b/vllm/transformers_utils/configs/k3_dspark.py new file mode 100644 index 0000000000..f1b21ed076 --- /dev/null +++ b/vllm/transformers_utils/configs/k3_dspark.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers import DeepseekV2Config + + +class K3DSparkConfig(DeepseekV2Config): + """Configuration for a dense MLA DSpark draft model.""" + + model_type = "k3_dspark" + has_no_defaults_at_init = True + + def __init__( + self, + mla_use_nope: bool = False, + mla_use_output_gate: bool = False, + mla_use_qk_norm: bool = False, + rope_theta: float = 50000.0, + **kwargs, + ) -> None: + # DeepseekV2Config defaults to a MoE topology. Zero these fields so + # generic vLLM config logic also recognizes this draft as dense. + kwargs.setdefault("n_routed_experts", 0) + kwargs.setdefault("n_shared_experts", 0) + kwargs.setdefault("num_experts_per_tok", 0) + + rope_parameters = kwargs.get("rope_parameters") + if rope_parameters is None: + kwargs["rope_parameters"] = { + "rope_type": "default", + "rope_theta": rope_theta, + } + else: + rope_parameters = dict(rope_parameters) + rope_parameters.setdefault("rope_type", "default") + rope_parameters.setdefault("rope_theta", rope_theta) + kwargs["rope_parameters"] = rope_parameters + + super().__init__(**kwargs) + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + self.mla_use_qk_norm = mla_use_qk_norm + + unsupported = [ + name + for name in ( + "mla_use_nope", + "mla_use_output_gate", + "mla_use_qk_norm", + "dspark_bonus_anchor", + ) + if getattr(self, name, False) + ] + if self.q_lora_rank is None: + unsupported.append("q_lora_rank=None") + if unsupported: + raise ValueError( + "MLA DSpark does not support " + ", ".join(unsupported) + "." + ) + + self.draft_vocab_size = ( + getattr(self, "draft_vocab_size", None) or self.vocab_size + ) + if self.draft_vocab_size != self.vocab_size: + raise ValueError( + "MLA DSpark requires draft_vocab_size to equal vocab_size when " + "sharing the target embedding and LM head." + ) + + target_layer_ids = getattr(self, "target_layer_ids", None) + if not target_layer_ids or getattr(self, "num_target_layers", None) != len( + target_layer_ids + ): + raise ValueError( + "MLA DSpark requires non-empty target_layer_ids and a matching " + "num_target_layers." + ) diff --git a/vllm/transformers_utils/configs/kimi_linear.py b/vllm/transformers_utils/configs/kimi_linear.py index 1489481680..6c2fd335ea 100644 --- a/vllm/transformers_utils/configs/kimi_linear.py +++ b/vllm/transformers_utils/configs/kimi_linear.py @@ -3,10 +3,6 @@ from transformers.configuration_utils import PretrainedConfig -from vllm.logger import init_logger - -logger = init_logger(__name__) - class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" @@ -49,8 +45,16 @@ class KimiLinearConfig(PretrainedConfig): qk_rope_head_dim: int | None = None, v_head_dim: int | None = None, mla_use_nope: bool | None = False, + mla_use_output_gate: bool | None = False, num_nextn_predict_layers: int = 0, linear_attn_config: dict | None = None, + attn_res_block_size: int | None = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: float | None = None, + activation_situ_linear_beta: float | None = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: int | None = None, + topk_method: str = "noaux_tc", **kwargs, ): self.model_type = model_type @@ -86,6 +90,7 @@ class KimiLinearConfig(PretrainedConfig): self.qk_rope_head_dim = qk_rope_head_dim self.v_head_dim = v_head_dim self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate # moe config self.num_experts = num_experts self.num_experts_per_token = num_experts_per_token @@ -102,6 +107,14 @@ class KimiLinearConfig(PretrainedConfig): self.topk_group = topk_group self.num_nextn_predict_layers = num_nextn_predict_layers + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + if linear_attn_config is not None: assert linear_attn_config["kda_layers"] is not None assert linear_attn_config["full_attn_layers"] is not None diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index e62b29295f..90c026cf9b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -267,6 +267,7 @@ class ModelArchConfigConvertorBase: "deepseek_v32", "deepseek_v4", "deepseek_mtp", + "k3_dspark", "glm_moe_dsa", "glm4_moe_lite", "glm4_moe_lite_mtp", diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index 60b56b75a0..4fa7d2fee8 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -27,6 +27,7 @@ __all__ = [ "IsaacProcessor", "KimiAudioProcessor", "KimiK25Processor", + "KimiK3Processor", "MiMoOmniProcessor", "MiniCPMOProcessor", "MiniCPMVProcessor", @@ -65,6 +66,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "IsaacProcessor": "vllm.transformers_utils.processors.isaac", "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio", "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25", + "KimiK3Processor": "vllm.transformers_utils.processors.kimi_k3", "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni", "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo", "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv", diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index 2203221c5a..5e4aec9b3c 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -27,8 +27,9 @@ def maybe_execute_in_parallel( """Run two functions potentially in parallel on separate CUDA streams. When aux_stream is provided, fn0 runs on the current (default) stream and - fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is - None, both functions execute sequentially on the current stream. + fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is + None or a breakable CUDA graph capture is active, both functions execute + sequentially on the current stream. This design follows TensorRT-LLM's maybe_execute_in_parallel pattern (tensorrt_llm/_torch/modules/multi_stream_utils.py). @@ -39,11 +40,18 @@ def maybe_execute_in_parallel( event0: CUDA event recorded before fn0 so aux_stream can wait. event1: CUDA event recorded after fn1 so default stream can wait. aux_stream: The second CUDA stream for fn1. - Multi-stream is disabled when aux_stream is None. + Multi-stream is disabled when aux_stream is None or a breakable + CUDA graph capture is active. Returns: Tuple of (fn0_result, fn1_result). """ + if aux_stream is not None: + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + if BreakableCUDAGraphCapture.is_active(): + aux_stream = None + if aux_stream is not None: event0.record() result0 = fn0() diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index 0a6eb14e0c..acd3b81816 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -5,7 +5,12 @@ from typing import ClassVar import torch from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla +from flashinfer.utils import ( + get_device_sm_count, + get_trtllm_gen_multi_ctas_kv_counter_bytes, +) +from vllm.config import get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( @@ -27,6 +32,37 @@ from vllm.v1.attention.backends.utils import KVCacheLayoutType logger = init_logger(__name__) + +def _trtllm_gen_mla_decode_supports_num_heads(num_heads: int) -> bool: + """True if trtllm-gen's MLA decode kernel supports this query head count. + + The kernel groups Q heads into CTAs of ``min(num_heads, tileSizeQ)`` and + requires ``num_heads`` divisible by that, else raises "The + numHeadsQ/numHeadsKv is not supported" (flashinfer fmhaKernels.cuh). + ``tileSizeQ`` is 8/16 for ``num_heads <= 8``/``<= 32`` (SwapsMmaAb) else 64; + treat ``> 32`` as tile 64 (safe bound). E.g. 96/24 -> False, 48/64/128 -> True. + """ + if num_heads <= 8: + tile = 8 + elif num_heads <= 32: + tile = 16 + else: + tile = 64 + return num_heads % min(num_heads, tile) == 0 + + +def _select_mla_decode_backend(num_heads: int) -> str | None: + """cute-dsl for head counts trtllm-gen cannot tile, else None (=> auto).""" + if not _trtllm_gen_mla_decode_supports_num_heads(num_heads): + logger.warning_once( + "trtllm-gen MLA decode does not support num_heads=%d " + "(query/kv head ratio); falling back to the cute-dsl backend.", + num_heads, + ) + return "cute-dsl" + return None + + FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024 @@ -48,9 +84,39 @@ def _get_workspace_buffer(return_lse: bool) -> torch.Tensor: return _fi_workspace +_fi_multi_ctas_kv_counter: torch.Tensor | None = None + + +def _get_multi_ctas_kv_counter_buffer( + min_bytes: int, device: torch.device +) -> torch.Tensor: + """Persistent, zero-initialized trtllm-gen multi-CTA-KV counter buffer. + + trtllm-gen's multi-CTA-KV MLA decode kernel resets these semaphores to zero + at the end of every launch, so the buffer only needs zeroing once. The + public ``trtllm_batch_decode_with_kv_cache_mla`` entry point builds a fresh + runner per call, so without a caller-owned buffer it re-allocates and + re-zeros this counter on every decode step (a tiny ``FillFunctor`` + launch right before the FMHA). Owning it here and passing it in removes that + per-step launch. ``min_bytes`` is sized to the worst-case batch so the + buffer is allocated once and never reallocated after CUDA-graph capture. + """ + global _fi_multi_ctas_kv_counter + if ( + _fi_multi_ctas_kv_counter is None + or _fi_multi_ctas_kv_counter.numel() < min_bytes + ): + _fi_multi_ctas_kv_counter = torch.zeros( + min_bytes, dtype=torch.uint8, device=device + ) + return _fi_multi_ctas_kv_counter + + class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. + supports_non_causal_multi_token_decode: ClassVar[bool] = True class FlashInferMLABackend(MLACommonBackend): @@ -91,6 +157,10 @@ class FlashInferMLABackend(MLACommonBackend): def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return capability.major == 10 + @classmethod + def supports_non_causal(cls) -> bool: + return True + @classmethod def supports_combination( cls, @@ -179,6 +249,16 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): self.bmm1_scale: float | None = None self.bmm2_scale: float | None = None + # Worst-case decode batch for the persistent trtllm-gen multi-CTA-KV + # counter buffer (see _get_multi_ctas_kv_counter_buffer). Captured here + # (config is in scope during construction) so the byte size can be + # resolved once on the first decode and never grows after CUDA-graph + # capture. The impl has no _vllm_config, hence get_current_vllm_config(). + _sched = get_current_vllm_config().scheduler_config + self._mla_counter_max_batch: int = ( + _sched.max_num_batched_tokens or _sched.max_num_seqs + ) + self._mla_counter_bytes: int | None = None def forward_mqa( self, @@ -194,8 +274,20 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): q_nope, q_pe = q q = torch.cat([q_nope, q_pe], dim=-1) + block_table = attn_metadata.decode.block_table + seq_lens = attn_metadata.decode.seq_lens + + if not attn_metadata.causal: + # Non-causal DSpark block: flatten to single-token decode rows with + # per-row context seq_lens (trtllm-gen has no causal flag and would + # otherwise mask the block causally). + query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes + q = q.unsqueeze(1) + if query_len > 1: + block_table = block_table.repeat_interleave(query_len, dim=0) + seq_lens = seq_lens.repeat_interleave(query_len) # trtllm API requires extra dimension q_len_per_request for MTP - if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: + elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: logger.warning_once( """FlashInferMLAImpl got a query of uneven length. This usually indicates an issue in batch reordering @@ -217,6 +309,29 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): return_lse = self.need_to_return_lse_for_decode workspace_buffer = _get_workspace_buffer(return_lse) + # trtllm-gen rejects MLA head counts it can't tile (e.g. 96); + # fall back to cute-dsl for those. + decode_backend = _select_mla_decode_backend(self.num_heads) + extra_kwargs = {} + if decode_backend: + extra_kwargs["backend"] = decode_backend + elif kv_c_and_k_pe_cache.shape[-2] in (32, 64): + # The auto path can dispatch to trtllm-gen, whose multi-CTA-KV decode + # kernel self-resets its semaphore counter after each launch (so it + # only needs zeroing once). Pass a persistent counter buffer to skip + # the per-step re-allocate + re-zero the public entry point would + # otherwise do. Guarded to configs where a trtllm-gen runner is + # eligible (page/block size in {32, 64}); the arg is rejected when + # only a cute-dsl runner can run. + if self._mla_counter_bytes is None: + self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes( + self._mla_counter_max_batch, + self.num_heads, + get_device_sm_count(q.device), + ) + extra_kwargs["multi_ctas_kv_counter_buffer"] = ( + _get_multi_ctas_kv_counter_buffer(self._mla_counter_bytes, q.device) + ) kernel_out = trtllm_batch_decode_with_kv_cache_mla( query=q, kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), @@ -224,12 +339,13 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): qk_nope_head_dim=self.qk_nope_head_dim, kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=attn_metadata.decode.block_table, - seq_lens=attn_metadata.decode.seq_lens, + block_tables=block_table, + seq_lens=seq_lens, max_seq_len=attn_metadata.max_seq_len, bmm1_scale=self.bmm1_scale, bmm2_scale=self.bmm2_scale, return_lse=return_lse, + **extra_kwargs, ) if return_lse: o, lse = kernel_out diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index dfcc037c8b..a516ee845c 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -73,6 +73,17 @@ class MLAPrefillBackend(ABC): pass. Overridden by backends that support it.""" return False + def supports_out(self) -> bool: + """Whether `run_prefill_new_tokens` honors a caller-provided `out` + tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the + final result into it in place. + + When True, callers may pass `out` and skip the post-hoc + slice/flatten/copy. False for backends that ignore `out` or emit a + padded (`qk_head_dim`) output. Overridden by backends that support it. + """ + return False + @classmethod def validate_configuration( cls, diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index d073f9b828..50897d3393 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -419,6 +419,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): return attn_out, lse return attn_out + def supports_out(self) -> bool: + # A padded V produces a qk_head_dim output that cannot be written into + # a v_head_dim `out`; only the unpadded path honors `out`. + return not self.requires_v_padding + def run_prefill_new_tokens( self, q: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 557c16f97f..eb1aef08bc 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -108,15 +108,21 @@ class FlashInferPrefillBackend(MLAPrefillBackend): return self._global_hyperparameters from vllm.model_executor.layers.attention.mla_attention import ( - MLAAttention, MLACommonImpl, ) + from vllm.model_executor.layers.attention_layer_base import ( + AttentionLayerBase, + ) + # Match any layer with an MLA impl, not just the MLAAttention wrapper: + # fused MLA modules (Kimi-K3's MultiHeadLatentAttention) register a + # different layer type. Keying on impl also excludes linear/KDA layers. forward_context = self.vllm_config.compilation_config.static_forward_context layer_names = [ name for name, layer in forward_context.items() - if isinstance(layer, MLAAttention) + if isinstance(layer, AttentionLayerBase) + and isinstance(getattr(layer, "impl", None), MLACommonImpl) ] self._global_hyperparameters = infer_global_hyperparameters( @@ -191,6 +197,10 @@ class FlashInferPrefillBackend(MLAPrefillBackend): o_data_type=prefill_metadata.output_dtype, ) + def supports_out(self) -> bool: + # Planned with head_dim_vo == v_head_dim, so the output is unpadded. + return True + def run_prefill_new_tokens( self, q: torch.Tensor, @@ -206,12 +216,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend): q=q, k=k, v=v, + out=out, return_lse=return_softmax_lse, ) if isinstance(ret, tuple): # Convert from (q_len, num_heads) to (num_heads, q_len) - return ret[0], ret[1].transpose(0, 1).contiguous() + return ret[0], ret[1].transpose(0, 1) return ret def run_prefill_context_chunk( @@ -229,4 +240,4 @@ class FlashInferPrefillBackend(MLAPrefillBackend): ) # Convert from (q_len, num_heads) to (num_heads, q_len) - return attn_out, lse.transpose(0, 1).contiguous() + return attn_out, lse.transpose(0, 1) diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 1f041f3731..780d9ae7f0 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -118,6 +118,10 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1] ) + def supports_out(self) -> bool: + # Output head dim is v_head_dim (unpadded); `out` supported since 0.1.8. + return True + def run_prefill_new_tokens( self, q: torch.Tensor, @@ -129,11 +133,11 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: from tokenspeed_mla import tokenspeed_mla_prefill - # `v` arrives as the second half of `kv_nope.split(...)` in - # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along - # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which - # would silently copy on a non-contiguous tensor; force contiguity here - # so the copy (if any) happens once outside the kernel call. + # `v` arrives as the second half of `kv_nope.split(...)` — a + # non-contiguous view of `kv_nope` along dim=-1. The kernel wraps inputs + # via `from_dlpack`, which preserves strides, but the FMHA kernel reads + # `v` assuming contiguous storage and silently produces wrong output on + # a strided view (verified on tokenspeed-mla 0.1.8). Force contiguity. v = v.contiguous() ret = tokenspeed_mla_prefill( @@ -148,11 +152,12 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): is_causal=True, return_lse=return_softmax_lse, enable_pdl=False, + out=out, ) if isinstance(ret, tuple): # Convert from (q_len, num_heads) to (num_heads, q_len) - return ret[0], ret[1].transpose(0, 1).contiguous() + return ret[0], ret[1].transpose(0, 1) return ret def run_prefill_context_chunk( @@ -188,4 +193,4 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): ) # Convert from (q_len, num_heads) to (num_heads, q_len) - return attn_out, lse.transpose(0, 1).contiguous() + return attn_out, lse.transpose(0, 1) diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 90f721272d..aa9a77dc42 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -91,6 +91,10 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1] ) + def supports_out(self) -> bool: + # Output head dim is v.shape[-1] == v_head_dim, so `out` is unpadded. + return True + def run_prefill_new_tokens( self, q: torch.Tensor, @@ -102,13 +106,14 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: from flashinfer.prefill import trtllm_ragged_attention_deepseek - out = torch.empty( - q.shape[0], - q.shape[1], - v.shape[2], - device=q.device, - dtype=self._prefill_metadata.output_dtype, - ) + if out is None: + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=self._prefill_metadata.output_dtype, + ) ret = trtllm_ragged_attention_deepseek( query=q, @@ -133,7 +138,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): if isinstance(ret, tuple): # Convert from (q_len, num_heads) to (num_heads, q_len) - return ret[0], ret[1].transpose(0, 1).contiguous() + return ret[0], ret[1].transpose(0, 1) return ret def run_prefill_context_chunk( @@ -182,4 +187,4 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): ) # Convert from (q_len, num_heads) to (num_heads, q_len) - return attn_out, lse.transpose(0, 1).contiguous() + return attn_out, lse.transpose(0, 1) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 977a42c6fe..4c7a90ed02 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -31,6 +31,32 @@ from vllm.v1.kv_cache_interface import AttentionSpec logger = init_logger(__name__) +@functools.lru_cache(maxsize=1) +def _get_mla_gluon(): + """Load the small-head Gluon MLA entry point.""" + unified_module = "aiter.ops.triton.gluon.mla_gluon" + try: + from aiter.ops.triton.gluon.mla_gluon import mla_gluon + + return mla_gluon + except ModuleNotFoundError as unified_import_error: + if not unified_module.startswith(unified_import_error.name or ""): + raise + legacy_module = "aiter.ops.triton.gluon.mla_decode_gluon" + try: + from aiter.ops.triton.gluon.mla_decode_gluon import mla_decode_gluon + + return mla_decode_gluon + except ModuleNotFoundError as legacy_import_error: + if not legacy_module.startswith(legacy_import_error.name or ""): + raise + raise RuntimeError( + "ROCM_AITER_MLA requires an AITER build with the small-head " + "Gluon MLA kernel (mla_gluon or mla_decode_gluon) when decode " + "heads are fewer than 16." + ) from unified_import_error + + @functools.lru_cache(maxsize=1) def _fp8_mla_prefill_supported() -> bool: """Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1. @@ -102,6 +128,10 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): attn_out_dtype: torch.dtype = torch.bfloat16 # The max query output length: int max_qo_len: int | None = None + # Minimum KV length used by Gluon to choose a safe split count. + min_kv_seq_len: int = 1 + # Small-head decode uses Gluon (avoids padding to 16). + use_gluon_decode: bool = False # Whether persistent MLA metadata was computed (only for qseqlen=1) has_persistent_metadata: bool = False @@ -241,7 +271,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): device=device, ) - self._fp8_prefill_enabled = _fp8_mla_prefill_supported() + # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads. + self._fp8_prefill_enabled = ( + _fp8_mla_prefill_supported() and self.num_heads % 16 == 0 + ) if self._fp8_prefill_enabled: max_prefill_qlen = min( vllm_config.model_config.max_model_len, @@ -469,6 +502,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): ) qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] max_qo_len = qo_len.max().item() + use_gluon_decode = AiterMLAHelper.use_gluon_decode( + self.num_heads, int(max_qo_len) + ) if self.compilation_config.cudagraph_mode.has_full_cudagraphs(): self.paged_kv_indices.fill_(-1) @@ -518,7 +554,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # so forward_mqa can skip passing it (falling back to the kernel # computing its own metadata internally, like v0.18.0). has_persistent_metadata = False - if max_qo_len == 1: + if max_qo_len == 1 and not use_gluon_decode: from aiter import get_mla_metadata_v1 get_mla_metadata_v1( @@ -553,6 +589,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): qo_indptr=qo_indptr, dcp_tot_seq_lens=dcp_tot_seq_lens_device, max_qo_len=max_qo_len, + use_gluon_decode=use_gluon_decode, attn_out_dtype=self.decode_attn_out_dtype, has_persistent_metadata=has_persistent_metadata, ) @@ -646,17 +683,16 @@ class AiterMLAHelper: @staticmethod def check_num_heads_validity(num_heads: int): assert AiterMLAHelper.is_valid_num_heads(num_heads), ( - f"Aiter MLA requires that num_heads be multiples or divisors of 16, " - f"but provided {num_heads} number of heads.\n" + "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple " + f"of 16 heads for persistent decode, but got {num_heads}.\n" f"Try adjusting tensor_parallel_size value." ) @staticmethod def is_valid_num_heads(num_heads: int) -> bool: - return ( - num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0 - if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS - else AiterMLAHelper._AITER_MIN_MLA_HEADS % num_heads == 0 + return num_heads > 0 and ( + num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS + or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0 ) @staticmethod @@ -681,6 +717,10 @@ class AiterMLAHelper: else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :] ) + @staticmethod + def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool: + return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1 + class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): def __init__( @@ -726,7 +766,10 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): # FP8 MLA prefill kernel imports (lazy, only when enabled). # Auto-enabled on gfx950 when AITER ships the kernels. - self._fp8_prefill_enabled = _fp8_mla_prefill_supported() + # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads. + self._fp8_prefill_enabled = ( + _fp8_mla_prefill_supported() and self.num_heads % 16 == 0 + ) if self._fp8_prefill_enabled: from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1 @@ -917,7 +960,106 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): ) -> tuple[torch.Tensor, torch.Tensor | None]: assert kv_c_and_k_pe_cache.numel() > 0 assert attn_metadata.decode is not None - assert attn_metadata.decode.max_qo_len is not None + + decode = attn_metadata.decode + assert decode.max_qo_len is not None + assert decode.paged_kv_indptr is not None + assert decode.paged_kv_indices is not None + if decode.use_gluon_decode: + if type(q) is tuple: + q_nope, q_pe = q + else: + q_nope, q_pe = torch.split( + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + B, num_q_heads, _ = q_nope.shape + o = torch.empty( + B, + num_q_heads, + self.kv_lora_rank, + dtype=decode.attn_out_dtype, + device=q_nope.device, + ) + kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1]) + mla_gluon = _get_mla_gluon() + mla_gluon( + q_nope=q_nope, + q_pe=q_pe, + kv_c=kv_buffer, + o=o, + page_table=decode.paged_kv_indices, + seq_info=decode.paged_kv_indptr, + sm_scale=self.scale, + k_pe=None, + kv_pe_offset=self.kv_lora_rank, + use_2d_view=False, + kv_scale=1.0, + min_kv_seq_len=decode.min_kv_seq_len, + ) + return o, None + + # 12-head (<16) non-causal multi-token verify (DSpark): the asm path has + # no gqa<16, qseqlen>1 kernel. Flatten each verify token to a qseqlen=1 + # gluon decode over the same committed prefix (non-causal), mirroring the + # TRITON_MLA / sparse-backend flatten but on the fast gluon kernel. Each + # request's KV metadata is repeated max_qo_len times. + if ( + self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS + and int(decode.max_qo_len) > 1 + ): + qlen = int(decode.max_qo_len) + if type(q) is tuple: + q_nope, q_pe = q + else: + q_nope, q_pe = torch.split( + q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + B, num_q_heads, _ = q_nope.shape + o = torch.empty( + B, + num_q_heads, + self.kv_lora_rank, + dtype=decode.attn_out_dtype, + device=q_nope.device, + ) + kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1]) + # Expand per-request paged-KV to per-verify-token: each request's KV + # block is repeated qlen times (row r*qlen+t attends request r's + # committed prefix). Fully vectorized (no host loop). + old_indptr = decode.paged_kv_indptr + per_req_len = old_indptr[1:] - old_indptr[:-1] + dev = q_nope.device + row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave( + qlen + ) + row_len = per_req_len[row_req] + new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to( + torch.int32 + ) + total = int(new_indptr[-1].item()) + within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[ + :-1 + ].to(torch.int64).repeat_interleave(row_len) + src = ( + old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within + ) + new_indices = decode.paged_kv_indices[src] + mla_gluon = _get_mla_gluon() + mla_gluon( + q_nope=q_nope, + q_pe=q_pe, + kv_c=kv_buffer, + o=o, + page_table=new_indices, + seq_info=new_indptr, + sm_scale=self.scale, + k_pe=None, + kv_pe_offset=self.kv_lora_rank, + use_2d_view=False, + kv_scale=1.0, + min_kv_seq_len=int(per_req_len.min()), + ) + return o, None if type(q) is tuple: q = torch.cat(q, dim=-1) @@ -960,11 +1102,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): kv_buffer, o, self.scale, - attn_metadata.decode.qo_indptr, - attn_metadata.decode.max_qo_len, - attn_metadata.decode.paged_kv_indptr, - attn_metadata.decode.paged_kv_indices, - attn_metadata.decode.paged_kv_last_page_len, + decode.qo_indptr, + decode.max_qo_len, + decode.paged_kv_indptr, + decode.paged_kv_indices, + decode.paged_kv_last_page_len, **mla_kwargs, ) diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 5c10c16a78..71a38e95d0 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -51,9 +51,17 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = ( AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE ) + # Non-causal DSpark block is flattened to one decode row per query token in + # forward_mqa, so no intra-block causal masking is required. + supports_non_causal_multi_token_decode: ClassVar[bool] = True def __init__(self, kv_cache_spec, layer_names, vllm_config, device): super().__init__(kv_cache_spec, layer_names, vllm_config, device) + # Only the non-causal DSpark draft group serves multi-token blocks via + # the decode path; raise its reorder threshold to the spec block length + # so full-cudagraph capture admits it. Causal usage stays single-token. + if getattr(self, "non_causal_multi_token_decode", False): + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self._reserve_attn_logits_workspace() def _reserve_attn_logits_workspace(self) -> None: @@ -68,6 +76,10 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): return # Decode reorder threshold is 1, so decode tokens <= max_num_seqs. B = self.vllm_config.scheduler_config.max_num_seqs + # Non-causal DSpark draft flattens each request's block to query_len + # decode rows; cover max_num_seqs * block_len rows. + if getattr(self, "non_causal_multi_token_decode", False): + B *= self.reorder_batch_threshold # DCP all-gathers the query heads before forward_mqa. q_num_heads = self.num_heads * self.dcp_world_size max_splits = _compute_num_kv_splits( @@ -132,6 +144,13 @@ class TritonMLABackend(MLACommonBackend): def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return True + @classmethod + def supports_non_causal(cls) -> bool: + # DSpark non-causal blocks are flattened to single-token decode rows in + # TritonMLAImpl.forward_mqa (decode_attention_fwd has no causal flag / + # no intra-block masking). Enables the non-causal AMD MLA path. + return True + class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): can_return_lse_for_decode: bool = True @@ -264,6 +283,18 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank] PAGE_SIZE = kv_c_and_k_pe_cache.size(1) + block_table = attn_metadata.decode.block_table + seq_lens = attn_metadata.decode.seq_lens + if not attn_metadata.causal: + # Non-causal DSpark block: flatten to one decode row per query token. + # Each row attends to the same committed KV prefix (per-row seq_lens) + # and never to sibling block tokens = non-causal block semantics. + # Mirrors FlashInferMLA's non-causal path. + query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes + if query_len > 1: + block_table = block_table.repeat_interleave(query_len, dim=0) + seq_lens = seq_lens.repeat_interleave(query_len) + # Run MQA — always pass layer scales. When KV cache is # BF16 the kernel's `if dtype.is_fp8()` check is a no-op. decode_attention_fwd( @@ -272,8 +303,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): kv_c_cache, o, lse, - attn_metadata.decode.block_table, - attn_metadata.decode.seq_lens, + block_table, + seq_lens, attn_logits, num_kv_splits, self.scale, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 18e6b36693..9aea061097 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1209,6 +1209,27 @@ def _get_kv_cache_groups_uniform_page_size( for layer_name, layer_spec in kv_cache_spec.items(): same_type_layers[layer_spec].append(layer_name) + # Attempt to further merge same-type layers based on whether their KV + # cache specs can be merged, to minimize the group count. This benefits + # situations where specs share a block layout and differ only in a + # property it can reconcile (e.g. full attention layers differing only in + # sliding window / attention chunk size). + layer_buckets: list[list[str]] = [] + spec_buckets: list[list[KVCacheSpec]] = [] + for layer_spec, layer_names in same_type_layers.items(): + for names, specs in zip(layer_buckets, spec_buckets): + try: + # A raise means that the specs are incompatible. + type(specs[0]).merge([*specs, layer_spec]) + except (AssertionError, ValueError): + continue + names.extend(layer_names) + specs.append(layer_spec) + break + else: + layer_buckets.append(list(layer_names)) + spec_buckets.append([layer_spec]) + # Split each group into smaller groups, to make the number of layers in each # group identical. Add padding to the last group of each type if necessary. # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) @@ -1221,9 +1242,9 @@ def _get_kv_cache_groups_uniform_page_size( # is the minimum number of layers among all attention types. Need a better # strategy if we want to support more complex patterns (e.g., 20 full + 30 # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in same_type_layers.values()]) + min_num_layers = min([len(layers) for layers in layer_buckets]) group_size = min_num_layers - max_num_layers = max([len(layers) for layers in same_type_layers.values()]) + max_num_layers = max([len(layers) for layers in layer_buckets]) if max_num_layers < min_num_layers * 1.5: # If the number of layers is not much larger than the minimum number of # layers, use the maximum number of layers as the group size to avoid @@ -1234,7 +1255,7 @@ def _get_kv_cache_groups_uniform_page_size( # extra layers to one attention type. group_size = max_num_layers grouped_layers = [] - for layers in same_type_layers.values(): + for layers in layer_buckets: num_padding_layers = group_size - len(layers) % group_size if num_padding_layers != group_size: logger.warning( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fa4a8daea9..99a8b44041 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -368,10 +368,11 @@ class Scheduler(SchedulerInterface): ) -> int: """Clip a prefill chunk so it ends where Mamba state must be cached. - In "align" cache mode the SSM state is only materialized at chunk - ends, so chunk ends are steered onto cacheable positions: block - boundaries by default, plus mandatory early stops (the prompt's - partial-tail hash boundary, a detected shared-prefix junction). + In "align" cache mode reusable SSM states are materialized at block + boundaries, plus mandatory early stops (the prompt's partial-tail hash + boundary, a detected shared-prefix junction). If a block is larger + than the configured prefill chunk limit, intermediate chunks keep + private running state until they reach the next cacheable position. """ start = ( request.num_computed_tokens @@ -392,11 +393,17 @@ class Scheduler(SchedulerInterface): last_cache_position = max(last_cache_position - block_size, 0) end = start + num_new_tokens - # Until `last_cache_position`, chunk ends must land on block - # boundaries. May yield an empty chunk (budget cannot reach the next - # boundary); the caller then skips the request. + # Until `last_cache_position`, prefer chunks ending on block + # boundaries. When a block cannot fit in any configured prefill chunk, + # allow sub-block progress and re-align at the next reachable boundary. if end < last_cache_position: - end = end // block_size * block_size + max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end next_block_boundary = (start // block_size + 1) * block_size tail_boundary = ( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 164350c7fb..b7199bddcb 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -385,6 +385,8 @@ class MLAAttentionSpec(FullAttentionSpec): alignment: int | None = None # Default to None for no padding. compress_ratio: int = 1 # Default to 1 for no compression. model_version: str | None = None + # Marks draft groups that flatten a non-causal query block into decode rows. + non_causal_multi_token_decode: bool = False def __post_init__(self): super().__post_init__() @@ -434,7 +436,7 @@ class MLAAttentionSpec(FullAttentionSpec): "quantization method, compress ratio, model version, and KV block " "stride indexing." ) - return cls( + merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, head_size=specs[0].head_size, @@ -445,7 +447,17 @@ class MLAAttentionSpec(FullAttentionSpec): cache_dtype_str=cache_dtype_str_set.pop(), compress_ratio=compress_ratio_set.pop(), model_version=model_version_set.pop(), + non_causal_multi_token_decode=any( + spec.non_causal_multi_token_decode for spec in specs + ), ) + for spec in specs: + for f in fields(AttentionSpec): + assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( + "All attention layers in the same KV cache group must have " + "the same attention spec." + ) + return merged_spec @dataclass(frozen=True, kw_only=True) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index ad7c40154b..c02c3350d8 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1006,11 +1006,11 @@ class SpecDecodeBaseProposer: def model_returns_tuple(self) -> bool: if self.method == "mtp": - # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final- - # norm hidden, so its forward returns (logit_hidden, - # recycle_hidden). Other MTP families return a single tensor. - return "DeepSeekMTPModel" in ( - self.draft_model_config.hf_config.architectures or [] + # These models return separate hidden states for logits and for + # feedback into the next draft step. + architectures = self.draft_model_config.hf_config.architectures or [] + return bool( + {"DeepSeekMTPModel", "KimiK3MTPModel"}.intersection(architectures) ) return self.method not in ("mtp", "draft_model", "dflash") @@ -1386,7 +1386,10 @@ class SpecDecodeBaseProposer: self.model.config.image_token_index = ( target_model.config.vision_config.image_token_id ) - elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration": + elif self.get_model_name(target_model) in ( + "KimiK25ForConditionalGeneration", + "KimiK3ForConditionalGeneration", + ): self.model.config.image_token_index = ( target_model.config.media_placeholder_token_id ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 645469e611..8f11e595d6 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -344,13 +344,6 @@ class DFlashSpeculator(DraftModelSpeculator): hidden_states = last_hidden_states self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens]) - self._copy_request_inputs( - num_reqs, - input_batch.idx_mapping, - temperature, - seeds, - ) - if dummy_run and skip_attn_for_dummy_run: # Memory profiling path: block_tables / kv_cache_config are not initialized. # Since DFlash needs to build its own attention metadata, we must skip the @@ -385,11 +378,15 @@ class DFlashSpeculator(DraftModelSpeculator): self.sample_indices, self.sample_pos, self.sample_idx_mapping, + self.temperature, + self.seeds, input_batch, num_sampled, num_rejected, last_sampled, next_prefill_tokens, + temperature, + seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], self.parallel_drafting_token_id, @@ -482,6 +479,8 @@ def _prepare_dflash_inputs_kernel( out_sample_indices_ptr, out_sample_pos_ptr, out_sample_idx_mapping_ptr, + out_temperature_ptr, + out_seeds_ptr, # Inputs from target batch target_positions_ptr, target_query_start_loc_ptr, @@ -490,6 +489,9 @@ def _prepare_dflash_inputs_kernel( next_prefill_tokens_ptr, num_sampled_ptr, num_rejected_ptr, + # Sampling params + temperature_ptr, + seeds_ptr, # Block table for slot mapping lookup. block_table_ptr, block_table_stride, @@ -585,6 +587,12 @@ def _prepare_dflash_inputs_kernel( # reads up to (context + query), not just the count of accepted # tokens this step. tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req) + # Copy sampling state. + tl.store( + out_temperature_ptr + req_state_idx, + tl.load(temperature_ptr + req_state_idx), + ) + tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx)) if req_idx == num_reqs - 1: # Pad per-request buffers to max_num_reqs for CUDA graph safety. last_query_end = num_reqs * num_query_per_req @@ -626,6 +634,8 @@ def prepare_dflash_inputs( sample_indices: torch.Tensor, sample_pos: torch.Tensor, sample_idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, input_batch: InputBatch, # [num_reqs] num_sampled: torch.Tensor, @@ -635,6 +645,10 @@ def prepare_dflash_inputs( last_sampled: torch.Tensor, # [max_num_reqs] next_prefill_tokens: torch.Tensor, + # [max_num_reqs] + input_temperature: torch.Tensor, + # [max_num_reqs] + input_seeds: torch.Tensor, # [max_num_reqs, max_num_blocks] block_table: torch.Tensor, block_size: int, @@ -665,6 +679,8 @@ def prepare_dflash_inputs( sample_indices, sample_pos, sample_idx_mapping, + temperature, + seeds, input_batch.positions, input_batch.query_start_loc, input_batch.idx_mapping, @@ -672,6 +688,8 @@ def prepare_dflash_inputs( next_prefill_tokens, num_sampled, num_rejected, + input_temperature, + input_seeds, block_table, block_table.stride(0), parallel_drafting_token_id, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4c300340c0..5083087bf0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6935,6 +6935,10 @@ class GPUModelRunner( num_active_loras=desc.num_active_loras, profile_seq_lens=profile_seq_lens, ) + if num_warmups > 0: + # Warmups may use auxiliary streams. Ensure all of their work has + # completed before beginning CUDA graph capture. + torch.accelerator.synchronize() with ( profiler, torch.profiler.record_function(