diff --git a/tests/core/framework/hf_model_loader_test.cpp b/tests/core/framework/hf_model_loader_test.cpp index 3a4b6e33c1..2db70a36c2 100644 --- a/tests/core/framework/hf_model_loader_test.cpp +++ b/tests/core/framework/hf_model_loader_test.cpp @@ -245,6 +245,27 @@ TEST(HFModelLoaderTest, RecFactoryCreatesRecCausalLmInstance) { #if defined(USE_NPU) || defined(USE_MLU) #if defined(USE_NPU) +TEST(HFModelLoaderTest, Qwen3DSparkFieldsFromTorchConfig) { + auto loader = ModelRegistry::get_model_args_loader("qwen3"); + ASSERT_NE(loader, nullptr); + + JsonReader reader; + ASSERT_TRUE(reader.parse_text(R"json( + { + "model_type": "qwen3", + "markov_rank": 256, + "enable_confidence_head": true, + "confidence_head_with_markov": true + } + )json")); + + ModelArgs args; + ASSERT_TRUE(loader(reader, &args)); + EXPECT_EQ(args.markov_rank(), 256); + EXPECT_TRUE(args.enable_confidence_head()); + EXPECT_TRUE(args.confidence_head_with_markov()); +} + TEST(HFModelLoaderTest, DeepseekV4DSparkModelArgsFrom0731Config) { auto loader = ModelRegistry::get_model_args_loader("deepseek_v4"); ASSERT_NE(loader, nullptr); diff --git a/tests/python/test_dspark.py b/tests/python/test_dspark.py new file mode 100644 index 0000000000..a07f4b43bf --- /dev/null +++ b/tests/python/test_dspark.py @@ -0,0 +1,112 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import pytest +import torch + +from xllm.python.models.dspark import ( + DSparkConfidenceHead, + DSparkForCausalLMBase, + DSparkMarkovHead, +) + + +def _model(*, enable_confidence_head: bool = True) -> DSparkForCausalLMBase: + return DSparkForCausalLMBase( + vocab_size=4, + draft_vocab_size=4, + markov_rank=2, + hidden_size=3, + enable_confidence_head=enable_confidence_head, + confidence_head_with_markov=True, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + +def test_markov_bias_matches_embedding_projection() -> None: + head = DSparkMarkovHead(4, 4, 2, torch.float32, torch.device("cpu")) + with torch.no_grad(): + head.markov_w1.weight.copy_( + torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 2.0], + [-1.0, 1.0], + ] + ) + ) + head.markov_w2.weight.copy_( + torch.tensor( + [ + [1.0, 0.0], + [0.0, 1.0], + [1.0, 1.0], + [2.0, -1.0], + ] + ) + ) + + output = head.bias(torch.tensor([0, 2])) + + expected = head.markov_w1(torch.tensor([0, 2])) @ head.markov_w2.weight.T + torch.testing.assert_close(output, expected) + + +def test_confidence_head_supports_batched_markov_features() -> None: + head = DSparkConfidenceHead(3, 2, True, torch.device("cpu")) + with torch.no_grad(): + head.proj.weight.copy_(torch.tensor([[1.0, -1.0, 0.5, 2.0, -2.0]])) + head.proj.bias.copy_(torch.tensor([0.25])) + hidden = torch.tensor([[[1.0, 2.0, 3.0], [0.5, 0.0, -1.0]]]) + markov = torch.tensor([[[0.25, 0.5], [1.0, -1.0]]]) + + output = head(hidden, markov) + + expected = torch.sigmoid(head.proj(torch.cat((hidden, markov), dim=-1))).squeeze(-1) + torch.testing.assert_close(output, expected) + + +def test_base_preserves_dspark_checkpoint_names() -> None: + assert set(_model().state_dict()) == { + "markov_head.markov_w1.weight", + "markov_head.markov_w2.weight", + "confidence_head.proj.weight", + "confidence_head.proj.bias", + } + + +def test_base_forwards_confidence_for_batched_hidden() -> None: + model = _model() + hidden = torch.ones(1, 2, 3) + prev_matrix = torch.tensor([[0, 1]]) + + output = model.dspark_confidence_probs(hidden, prev_matrix) + + confidence_head = model.confidence_head + assert confidence_head is not None + expected = confidence_head(hidden, model.markov_head.embed(prev_matrix)) + torch.testing.assert_close(output, expected) + assert model.has_dspark_confidence_head() + + +def test_base_rejects_confidence_without_head() -> None: + model = _model(enable_confidence_head=False) + + with pytest.raises(RuntimeError, match="not enabled"): + model.dspark_confidence_probs(torch.ones(1, 1, 3), torch.tensor([[0]])) + assert not model.has_dspark_confidence_head() diff --git a/tests/python/test_glm5_2_parallel.py b/tests/python/test_glm5_2_parallel.py index 70de4e59f4..707529591a 100644 --- a/tests/python/test_glm5_2_parallel.py +++ b/tests/python/test_glm5_2_parallel.py @@ -74,7 +74,9 @@ def test_full_world_ep_partitions_glm_experts() -> None: assert moe.local_expert_start == 6 assert moe.local_expert_end == 8 + moe.allocate_experts_w13_for_loading() assert moe.experts_w13.shape == (2, 16, 16) + moe.allocate_experts_w2_for_loading() assert moe.experts_w2.shape == (2, 16, 8) diff --git a/tests/python/test_qwen3_capture.py b/tests/python/test_qwen3_capture.py new file mode 100644 index 0000000000..2feb4a2410 --- /dev/null +++ b/tests/python/test_qwen3_capture.py @@ -0,0 +1,127 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +import xllm.python.models.qwen3 as qwen3_module +from xllm.python.models.aux_hidden_capture import AuxHiddenCapture +from xllm.python.models.qwen3 import Qwen3Config, Qwen3Model + + +class _Embedding(nn.Module): + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + values = input_ids.to(torch.float32) + return torch.stack((values, values + 10.0), dim=-1) + + +class _ResidualLayer(nn.Module): + def __init__(self, delta: float) -> None: + super().__init__() + self.delta = delta + + def forward( + self, + hidden: torch.Tensor, + residual: torch.Tensor | None, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + cos: torch.Tensor | None, + sin: torch.Tensor | None, + mrope_section: list[int] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + del positions, cos_sin_cache, cos, sin, mrope_section + residual = hidden if residual is None else hidden + residual + return torch.full_like(hidden, self.delta), residual + + +class _FinalNorm(nn.Module): + def forward( + self, + hidden: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + return (hidden if residual is None else hidden + residual), residual + + +def _config(*, layers_to_capture: tuple[int, ...]) -> Qwen3Config: + return Qwen3Config( + hidden_size=2, + n_layers=3, + n_heads=1, + n_kv_heads=1, + head_dim=2, + intermediate_size=4, + max_position_embeddings=8, + vocab_size=4, + layers_to_capture=layers_to_capture, + ) + + +def _model(monkeypatch: pytest.MonkeyPatch, layers_to_capture: tuple[int, ...]) -> Qwen3Model: + monkeypatch.setattr( + qwen3_module, + "get_forward_context", + lambda: SimpleNamespace(cp_context=None), + ) + model = Qwen3Model(_config(layers_to_capture=layers_to_capture), torch.float32, torch.device("cpu")) + model.embed_tokens = _Embedding() + model.layers = nn.ModuleList([_ResidualLayer(1.0), _ResidualLayer(2.0), _ResidualLayer(3.0)]) + model.norm = _FinalNorm() + return model + + +def test_qwen3_config_reads_capture_layers() -> None: + config = Qwen3Config.from_dict({"layers_to_capture": [3, 1]}) + + assert config.layers_to_capture == (3, 1) + + +def test_qwen3_model_returns_captured_residual_streams_in_config_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _model(monkeypatch, layers_to_capture=(2, 1)) + embedded = model.embed_tokens(torch.tensor([1, 2])) + + output = model(torch.tensor([1, 2]), torch.tensor([0, 1])) + + assert isinstance(output, tuple) + hidden, aux_hidden = output + torch.testing.assert_close(hidden, embedded + 6.0) + torch.testing.assert_close(aux_hidden, torch.cat((embedded + 3.0, embedded + 1.0), dim=-1)) + + +def test_qwen3_model_returns_tensor_when_capture_is_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + model = _model(monkeypatch, layers_to_capture=()) + + output = model(torch.tensor([1, 2]), torch.tensor([0, 1])) + + assert isinstance(output, torch.Tensor) + + +def test_aux_hidden_capture_snapshots_hidden_without_residual() -> None: + capture = AuxHiddenCapture((0,)) + hidden = torch.tensor([[1.0, 2.0]]) + captured: dict[int, torch.Tensor] = {} + + capture.capture_layer(0, hidden, None, captured) + hidden.add_(10.0) + _, aux_hidden = capture.finalize(hidden, captured) + + torch.testing.assert_close(aux_hidden, torch.tensor([[1.0, 2.0]])) diff --git a/tests/python/test_qwen3_dflash.py b/tests/python/test_qwen3_dflash.py new file mode 100644 index 0000000000..606b9255cf --- /dev/null +++ b/tests/python/test_qwen3_dflash.py @@ -0,0 +1,244 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch + +from xllm.python import distributed, kernels +from xllm.python.models.qwen3_dflash import ( + DFlashContextProjection, + DFlashQwen3Config, + DFlashQwen3ForCausalLM, + DFlashQwen3Model, +) + + +def _config_dict(**overrides) -> dict: + values = { + "hidden_size": 4, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 8, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 16, + "vocab_size": 8, + "draft_vocab_size": 8, + "tp_size": 1, + "tp_rank": 0, + "dp_size": 1, + "dp_rank": 0, + "dtype": "float32", + "device": "cpu", + } + values.update(overrides) + return values + + +def _config(**overrides) -> DFlashQwen3Config: + config = DFlashQwen3Config.from_dict(_config_dict(**overrides)) + config.validate() + return config + + +class _StateDict: + def __init__(self, tensors: dict[str, torch.Tensor]) -> None: + self._tensors = tensors + + def has(self, name: str) -> bool: + return name in self._tensors + + def get_tensor(self, name: str) -> torch.Tensor: + return self._tensors[name] + + +def test_top_level_and_nested_rope_config_formats_are_supported() -> None: + top_level_rope = _config(model_type="qwen3", rope_theta=10000.0) + nested_rope = _config( + model_type="qwen3", + rope_theta=None, + rope_parameters={"rope_theta": 1e7}, + ) + + assert top_level_rope.rope_theta == 10000.0 + assert nested_rope.rope_theta == 1e7 + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"dp_size": 2, "world_size": 1}, "world_size must equal"), + ({"dp_size": 2, "dp_rank": 2, "world_size": 2}, "dp_rank must be"), + ({"tp_rank": 1}, "tp_rank must be"), + ], +) +def test_config_rejects_invalid_parallel_settings( + overrides: dict, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + _config(**overrides) + + +def test_config_rejects_reduced_draft_vocabulary() -> None: + with pytest.raises(ValueError, match="reduced-vocabulary"): + _config(draft_vocab_size=4) + + +def test_draft_attention_is_non_causal() -> None: + model = DFlashQwen3Model(_config(), torch.float32, torch.device("cpu")) + + assert not model.layers[0].self_attn.attn.causal + + +def test_context_projection_uses_tensor_parallel_output_shard( + monkeypatch: pytest.MonkeyPatch, +) -> None: + projection = DFlashContextProjection( + out_features=4, + tp_size=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + weight = torch.arange(12, dtype=torch.float32).view(4, 3) + hidden = torch.tensor([[1.0, 2.0, 3.0]]) + rank_zero_output = torch.nn.functional.linear(hidden, weight[:2]) + all_gather = Mock( + side_effect=lambda local_output, **_: torch.cat( + (rank_zero_output, local_output), + dim=-1, + ) + ) + monkeypatch.setattr(distributed, "all_gather", all_gather, raising=False) + + projection.load_weight(weight, tp_rank=1) + output = projection(hidden) + + torch.testing.assert_close(projection.weight, weight[2:]) + torch.testing.assert_close(output, torch.nn.functional.linear(hidden, weight)) + all_gather.assert_called_once() + torch.testing.assert_close( + all_gather.call_args.args[0], + torch.nn.functional.linear(hidden, weight[2:]), + ) + assert all_gather.call_args.kwargs == {"dim": -1, "world_size": 2} + + +def test_context_projection_writes_each_layer_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = _config() + model = DFlashQwen3Model(config, torch.float32, torch.device("cpu")) + attention = model.layers[0].self_attn + with torch.no_grad(): + model.fc.load_weight(torch.eye(config.hidden_size), tp_rank=0) + model.hidden_norm.weight.fill_(1.0) + attention.qkv_proj.weight.zero_() + attention.qkv_proj.weight[attention.q_size : attention.q_size + attention.kv_size].copy_( + torch.eye(config.hidden_size) + ) + attention.qkv_proj.weight[attention.q_size + attention.kv_size :].copy_(2.0 * torch.eye(config.hidden_size)) + attention.k_norm.weight.fill_(1.0) + model._build_context_kv_buffers() + monkeypatch.setattr( + kernels, + "rms_norm", + lambda hidden, weight, eps: hidden + * torch.rsqrt(hidden.float().pow(2).mean(dim=-1, keepdim=True) + eps) + * weight, + raising=False, + ) + reshape_paged_cache = Mock() + monkeypatch.setattr( + kernels, + "reshape_paged_cache", + reshape_paged_cache, + raising=False, + ) + synchronizer = Mock() + key_cache = torch.empty(1, 1, 1, config.head_dim) + value_cache = torch.empty_like(key_cache) + target_hidden = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + + projected = model.write_context_kv( + target_hidden, + torch.tensor([0]), + torch.tensor([0], dtype=torch.int32), + [(key_cache, value_cache, None, None, None)], + synchronizer, + ) + + reshape_paged_cache.assert_called_once() + call_args = reshape_paged_cache.call_args.args + torch.testing.assert_close( + call_args[2], + 2.0 * projected.view(1, 1, config.head_dim), + ) + assert call_args[3] is key_cache + assert call_args[4] is value_cache + synchronizer.record_event.assert_called_once_with(0) + + +def test_checkpoint_weight_names_load_into_fused_modules( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + kernels, + "prepare_row_parallel_weight", + lambda weight: (weight, False), + raising=False, + ) + model = DFlashQwen3ForCausalLM(_config_dict()) + tensors = { + "fc.weight": torch.ones(4, 8), + "hidden_norm.weight": torch.ones(4), + "layers.0.input_layernorm.weight": torch.ones(4), + "layers.0.post_attention_layernorm.weight": torch.ones(4), + "layers.0.self_attn.q_norm.weight": torch.ones(4), + "layers.0.self_attn.k_norm.weight": torch.ones(4), + "layers.0.self_attn.q_proj.weight": torch.full((4, 4), 1.0), + "layers.0.self_attn.k_proj.weight": torch.full((4, 4), 2.0), + "layers.0.self_attn.v_proj.weight": torch.full((4, 4), 3.0), + "layers.0.self_attn.o_proj.weight": torch.full((4, 4), 4.0), + "layers.0.mlp.gate_proj.weight": torch.full((8, 4), 5.0), + "layers.0.mlp.up_proj.weight": torch.full((8, 4), 6.0), + "layers.0.mlp.down_proj.weight": torch.full((4, 8), 7.0), + "norm.weight": torch.ones(4), + } + + model.load_weights([_StateDict(tensors)], tp_rank=0, tp_size=1) + + qkv_weight = model.model.layers[0].self_attn.qkv_proj.weight + torch.testing.assert_close( + qkv_weight[:4], + tensors["layers.0.self_attn.q_proj.weight"], + ) + torch.testing.assert_close( + qkv_weight[4:8], + tensors["layers.0.self_attn.k_proj.weight"], + ) + torch.testing.assert_close( + qkv_weight[8:12], + tensors["layers.0.self_attn.v_proj.weight"], + ) + assert model.model._fused_kv_weight.shape == (8, 4) + assert model.model.fc.weight.shape == (4, 8) + assert model.model.embed_tokens is None + assert model.lm_head is None diff --git a/tests/python/test_qwen3_dspark.py b/tests/python/test_qwen3_dspark.py new file mode 100644 index 0000000000..b08125890a --- /dev/null +++ b/tests/python/test_qwen3_dspark.py @@ -0,0 +1,127 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import pytest +import torch + +from xllm.python import kernels +from xllm.python.models.qwen3_dspark import ( + Qwen3DSparkConfig, + Qwen3DSparkForCausalLM, +) + + +def _config(**overrides) -> Qwen3DSparkConfig: + values = _config_dict(**overrides) + config = Qwen3DSparkConfig.from_dict(values) + config.validate() + return config + + +def _config_dict(**overrides) -> dict: + values = { + "hidden_size": 4, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 4, + "intermediate_size": 8, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 16, + "vocab_size": 8, + "draft_vocab_size": 8, + "markov_rank": 2, + "tp_size": 1, + "tp_rank": 0, + "dp_size": 1, + "dp_rank": 0, + "dtype": "float32", + "device": "cpu", + } + values.update(overrides) + return values + + +class _StateDict: + def __init__(self, tensors: dict[str, torch.Tensor]) -> None: + self._tensors = tensors + + def has(self, name: str) -> bool: + return name in self._tensors + + def get_tensor(self, name: str) -> torch.Tensor: + return self._tensors[name] + + +def test_config_requires_positive_markov_rank() -> None: + with pytest.raises(ValueError, match="markov_rank > 0"): + _config(markov_rank=0) + + +def test_checkpoint_loads_dspark_heads(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + kernels, + "prepare_row_parallel_weight", + lambda weight: (weight, False), + raising=False, + ) + model = Qwen3DSparkForCausalLM( + _config_dict( + enable_confidence_head=True, + confidence_head_with_markov=True, + ) + ) + tensors = { + "fc.weight": torch.eye(4), + "hidden_norm.weight": torch.ones(4), + "layers.0.input_layernorm.weight": torch.ones(4), + "layers.0.post_attention_layernorm.weight": torch.ones(4), + "layers.0.self_attn.q_norm.weight": torch.ones(4), + "layers.0.self_attn.k_norm.weight": torch.ones(4), + "layers.0.self_attn.q_proj.weight": torch.full((4, 4), 1.0), + "layers.0.self_attn.k_proj.weight": torch.full((4, 4), 2.0), + "layers.0.self_attn.v_proj.weight": torch.full((4, 4), 3.0), + "layers.0.self_attn.o_proj.weight": torch.full((4, 4), 4.0), + "layers.0.mlp.gate_proj.weight": torch.full((8, 4), 5.0), + "layers.0.mlp.up_proj.weight": torch.full((8, 4), 6.0), + "layers.0.mlp.down_proj.weight": torch.full((4, 8), 7.0), + "norm.weight": torch.ones(4), + "markov_head.markov_w1.weight": torch.full((8, 2), 8.0), + "markov_head.markov_w2.weight": torch.full((8, 2), 9.0), + "confidence_head.proj.weight": torch.full((1, 6), 10.0), + "confidence_head.proj.bias": torch.full((1,), 11.0), + } + + model.load_weights([_StateDict(tensors)], tp_rank=0, tp_size=1) + + torch.testing.assert_close( + model.markov_head.markov_w1.weight, + tensors["markov_head.markov_w1.weight"], + ) + torch.testing.assert_close( + model.markov_head.markov_w2.weight, + tensors["markov_head.markov_w2.weight"], + ) + assert model.confidence_head is not None + torch.testing.assert_close( + model.confidence_head.proj.weight, + tensors["confidence_head.proj.weight"], + ) + torch.testing.assert_close( + model.confidence_head.proj.bias, + tensors["confidence_head.proj.bias"], + ) diff --git a/tests/python/test_registry.py b/tests/python/test_registry.py index c61b534a03..999e9569b7 100644 --- a/tests/python/test_registry.py +++ b/tests/python/test_registry.py @@ -21,10 +21,10 @@ def test_unsupported_model_fails_before_import(monkeypatch: pytest.MonkeyPatch) -> None: import_model = Mock() - monkeypatch.setattr(registry.current_platform, "device_type", lambda: "npu") + monkeypatch.setattr(registry.current_platform, "device_type", lambda: "cuda") monkeypatch.setattr(registry, "import_module", import_model) - with pytest.raises(NotImplementedError, match="qwen3_5.*npu"): - registry.get_model_class("qwen3_5") + with pytest.raises(NotImplementedError, match="qwen3_vl.*cuda"): + registry.get_model_class("qwen3_vl") import_model.assert_not_called() diff --git a/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index acd5536383..9d1a1aa56b 100644 --- a/xllm/core/kernels/npu/npu_ops_library.cpp +++ b/xllm/core/kernels/npu/npu_ops_library.cpp @@ -32,6 +32,7 @@ limitations under the License. #include "kernels/npu/xllm_ops/xllm_ops_api.h" #include "npu_ops_api.h" +#include "triton_npu/torch_api/triton_ops_api.h" namespace xllm { @@ -43,6 +44,159 @@ torch::Tensor rms_norm_npu(const torch::Tensor& input, return xllm::kernel::npu::rms_norm(input, weight, eps, "rmsnorm"); } +torch::Tensor rms_norm_gated_npu(const torch::Tensor& input, + const torch::Tensor& gate, + const torch::Tensor& weight, + double eps) { + return xllm::kernel::npu::layer_norm_fwd_aclnn(input, + weight, + /*bias=*/torch::Tensor(), + eps, + /*z=*/gate, + /*group_size=*/input.size(-1), + /*norm_before_gate=*/true, + /*is_rms_norm=*/true); +} + +torch::Tensor l2_norm_npu(torch::Tensor input, double eps) { + return xllm::kernel::npu::npu_l2norm_last_dim(input, eps); +} + +torch::Tensor causal_conv1d_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + // Convert device tensors to host vectors for IntArrayRef parameters. + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + constexpr int64_t kActivationSilu = 1; + constexpr int64_t kPadSlotId = -1; + constexpr int64_t kRunModeForward = 0; + + return xllm::kernel::npu::causal_conv1d( + x, + weight, + conv_state, + /*bias_opt=*/std::nullopt, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + /*num_accepted_tokens_opt=*/torch::IntArrayRef{}, + kActivationSilu, + kPadSlotId, + kRunModeForward); +} + +std::tuple +causal_conv1d_qkv_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc, + int64_t num_qk_heads, + int64_t num_v_heads, + int64_t head_k_dim, + int64_t head_v_dim) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + return xllm::kernel::npu::causal_conv1d_qkv(x, + weight, + conv_state, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + num_qk_heads, + num_v_heads, + head_k_dim, + head_v_dim); +} + +std::tuple chunk_gated_delta_rule_npu( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor initial_state, + torch::Tensor cu_seqlens) { + return xllm::kernel::npu::npu_mega_chunk_gdn( + q, + k, + v, + g, + beta, + /*scale=*/std::nullopt, + /*initial_state=*/initial_state, + /*output_final_state=*/true, + /*cu_seqlens=*/cu_seqlens, + /*q_seq_lens=*/{}, + /*use_qk_l2norm_in_kernel=*/true); +} + +torch::Tensor fused_sigmoid_gating_delta_rule_decode_npu( + torch::Tensor a_log, + torch::Tensor a, + torch::Tensor dt_bias, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor b, + torch::Tensor ssm_state, + torch::Tensor state_indices, + torch::Tensor cu_seqlens, + double scale) { + auto a_log_f32 = a_log.to(torch::kFloat32); + auto dt_bias_f32 = dt_bias.to(torch::kFloat32); + return xllm::kernel::npu::npu_fused_sigmoid_gating_delta_rule_update( + a_log_f32, + a, + dt_bias_f32, + q, + k, + v, + b, + ssm_state, + state_indices, + cu_seqlens, + /*scale=*/static_cast(scale), + /*use_qk_l2norm_in_kernel=*/true, + /*softplus_beta=*/1.0f, + /*softplus_threshold=*/20.0f); +} + std::tuple fused_add_rms_norm_npu( torch::Tensor& input, torch::Tensor& residual, @@ -293,6 +447,29 @@ void ensure_xllm_ops_registered() { // compiled only under USE_NPU (mutually exclusive with USE_CUDA). TORCH_LIBRARY(xllm_ops, m) { m.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor"); + m.def( + "rms_norm_gated(Tensor input, Tensor gate, Tensor weight, float eps) -> " + "Tensor"); + m.def("l2_norm(Tensor input, float eps) -> Tensor"); + m.def( + "chunk_gated_delta_rule(Tensor q, Tensor k, Tensor v, Tensor g, " + "Tensor beta, Tensor initial_state, Tensor cu_seqlens) -> " + "(Tensor, Tensor)"); + m.def( + "causal_conv1d_prefill(Tensor x, Tensor weight, Tensor(a!) conv_state, " + "Tensor state_indices, Tensor has_initial_state, " + "Tensor query_start_loc) -> Tensor"); + m.def( + "causal_conv1d_qkv_prefill(Tensor x, Tensor weight, " + "Tensor(a!) conv_state, Tensor state_indices, " + "Tensor has_initial_state, Tensor query_start_loc, " + "int num_qk_heads, int num_v_heads, " + "int head_k_dim, int head_v_dim) -> (Tensor, Tensor, Tensor)"); + m.def( + "fused_sigmoid_gating_delta_rule_decode(Tensor a_log, Tensor a, " + "Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor b, " + "Tensor(a!) ssm_state, Tensor state_indices, Tensor cu_seqlens, " + "float scale) -> Tensor"); m.def( "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor " "weight, " @@ -398,6 +575,14 @@ TORCH_LIBRARY(xllm_ops, m) { TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("rms_norm", TORCH_FN(xllm::rms_norm_npu)); + m.impl("rms_norm_gated", TORCH_FN(xllm::rms_norm_gated_npu)); + m.impl("l2_norm", TORCH_FN(xllm::l2_norm_npu)); + m.impl("chunk_gated_delta_rule", TORCH_FN(xllm::chunk_gated_delta_rule_npu)); + m.impl("causal_conv1d_prefill", TORCH_FN(xllm::causal_conv1d_prefill_npu)); + m.impl("causal_conv1d_qkv_prefill", + TORCH_FN(xllm::causal_conv1d_qkv_prefill_npu)); + m.impl("fused_sigmoid_gating_delta_rule_decode", + TORCH_FN(xllm::fused_sigmoid_gating_delta_rule_decode_npu)); m.impl("fused_add_rms_norm", TORCH_FN(xllm::fused_add_rms_norm_npu)); m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu)); m.impl("inplace_partial_rotary_mul", diff --git a/xllm/core/runtime/dflash_worker_impl.cpp b/xllm/core/runtime/dflash_worker_impl.cpp index 320ddbb474..52c0a51c1c 100644 --- a/xllm/core/runtime/dflash_worker_impl.cpp +++ b/xllm/core/runtime/dflash_worker_impl.cpp @@ -388,17 +388,21 @@ bool DFlashWorkerImpl::init_model(const std::string& model_weights_path, // the draft backbone/Markov head project through the wrong vocabulary // basis and reduces acceptance to near-random levels. } else { + const bool python_weights_shared = + draft_impl_->share_weights_from(*impl_); + if (!python_weights_shared) { #if defined(USE_NPU) - auto head = impl_->get_npu_lm_head(); - draft_impl_->set_npu_lm_head(head); - auto word_embedding = impl_->get_npu_word_embedding(); - draft_impl_->set_npu_word_embedding(word_embedding); + auto head = impl_->get_npu_lm_head(); + draft_impl_->set_npu_lm_head(head); + auto word_embedding = impl_->get_npu_word_embedding(); + draft_impl_->set_npu_word_embedding(word_embedding); #else - auto head = impl_->get_lm_head(); - draft_impl_->set_lm_head(head); - auto word_embedding = impl_->get_word_embedding(); - draft_impl_->set_word_embedding(word_embedding); + auto head = impl_->get_lm_head(); + draft_impl_->set_lm_head(head); + auto word_embedding = impl_->get_word_embedding(); + draft_impl_->set_word_embedding(word_embedding); #endif + } } JsonReader reader; diff --git a/xllm/core/runtime/py_executor_impl.cpp b/xllm/core/runtime/py_executor_impl.cpp index 8cc46c9c47..5a5f828fa2 100644 --- a/xllm/core/runtime/py_executor_impl.cpp +++ b/xllm/core/runtime/py_executor_impl.cpp @@ -240,6 +240,14 @@ ModelOutput PyExecutorImpl::run(const torch::Tensor& tokens, // mRoPE [3,N]->1-D decode collapse. py::object hidden_obj = py_executor_.attr("execute")( tokens, positions_arg, py_metadata, input_embedding, py_sync); + if (py::isinstance(hidden_obj)) { + py::tuple output = hidden_obj.cast(); + CHECK_EQ(output.size(), 2) << "Python model tuple output must be " + "(hidden_states, aux_hidden_states)"; + return ModelOutput(output[0].cast(), + torch::Tensor(), + output[1].cast()); + } return ModelOutput(hidden_obj.cast()); } diff --git a/xllm/models/llm/npu/deepseek_v2.h b/xllm/models/llm/npu/deepseek_v2.h index b3b9a8a6a1..821697fee1 100644 --- a/xllm/models/llm/npu/deepseek_v2.h +++ b/xllm/models/llm/npu/deepseek_v2.h @@ -175,7 +175,8 @@ class DeepseekV2ModelImpl : public torch::nn::Module { if (::xllm::KVCacheConfig::get_instance().enable_prefix_cache() && !input_params.meta.batch_forward_type.is_decode()) { attn_mask = attn_mask_.get_attn_mask(512, dtype_, device_); - } else if (input_params.meta.batch_forward_type.is_prefill()) { + } else if (input_params.meta.batch_forward_type.is_prefill() || + input_params.meta.batch_forward_type.is_chunked_prefill()) { attn_mask = attn_mask_.get_attn_mask(128, dtype_, device_); } else if (num_speculative_tokens_ > 0) { // TODO :the judgement of gen_free_mask need more check diff --git a/xllm/models/llm/py_causal_lm.cpp b/xllm/models/llm/py_causal_lm.cpp index cedfbdc5da..e26fa52678 100644 --- a/xllm/models/llm/py_causal_lm.cpp +++ b/xllm/models/llm/py_causal_lm.cpp @@ -30,6 +30,10 @@ limitations under the License. #include "core/framework/state_dict/state_dict.h" #include "models/py_model_helper.h" +#if defined(USE_NPU) +#include "platform/npu/npu_layer_synchronizer.h" +#endif + namespace py = pybind11; namespace xllm { @@ -47,6 +51,23 @@ void share_python_model_weights(py::object& draft_model, namespace { +py::object optional_tensor(const torch::Tensor& tensor) { + return tensor.defined() ? py::cast(tensor) : py::none(); +} + +py::list build_python_kv_caches(std::vector& kv_caches) { + py::list python_caches; + for (KVCache& kv_cache : kv_caches) { + python_caches.append( + py::make_tuple(optional_tensor(kv_cache.get_k_cache()), + optional_tensor(kv_cache.get_v_cache()), + optional_tensor(kv_cache.get_index_cache()), + optional_tensor(kv_cache.get_conv_cache()), + optional_tensor(kv_cache.get_ssm_cache()))); + } + return python_caches; +} + void clear_python_object(py::object& object) { if (!object) { return; @@ -182,10 +203,24 @@ PyCausalLM::PyCausalLM(const ModelContext& context) } PyCausalLM::~PyCausalLM() { + clear_python_object(python_kv_caches_); clear_python_object(py_model_); clear_python_object(config_dict_); } +const py::object& PyCausalLM::get_or_build_python_kv_caches( + std::vector& kv_caches) { + const int64_t num_layers = static_cast(kv_caches.size()); + if (!python_kv_caches_) { + python_kv_caches_ = build_python_kv_caches(kv_caches); + python_kv_cache_layer_count_ = num_layers; + } else { + CHECK_EQ(num_layers, python_kv_cache_layer_count_) + << "KV cache layer count changed after initial Python conversion"; + } + return python_kv_caches_; +} + py::dict PyCausalLM::build_config_dict( const ParallelArgs& parallel_args) const { py::dict d; @@ -205,9 +240,17 @@ py::dict PyCausalLM::build_config_dict( // cp_size is a reflected ParallelArgs PROPERTY (already in d), but cp_rank is // a derived member function, so pass it explicitly for the Python executor. d["cp_rank"] = cp_rank_; - d["enable_graph"] = ExecutionConfig::get_instance().enable_graph(); + const bool requires_eager_execution = + !model_args_.layers_to_capture().empty() || + model_args_.model_type() == "DFlashDraftModel" || + model_args_.model_type() == "DSparkDraftModel"; + d["enable_graph"] = requires_eager_execution + ? false + : ExecutionConfig::get_instance().enable_graph(); d["python_graph_backend"] = - ExecutionConfig::get_instance().python_graph_backend(); + requires_eager_execution + ? std::string("off") + : ExecutionConfig::get_instance().python_graph_backend(); return d; } @@ -247,6 +290,57 @@ torch::Tensor PyCausalLM::logits(const torch::Tensor& hidden_states, return out.cast(); } +ModelOutput PyCausalLM::write_context_kv( + const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) { + torch::NoGradGuard no_grad; + py::gil_scoped_acquire gil; + py::object layer_synchronizer = py::none(); +#if defined(USE_NPU) + if (input_params.parallel.layer_synchronizer != nullptr) { + layer_synchronizer = py::cast(input_params.parallel.layer_synchronizer); + } +#endif + const py::object& python_kv_caches = get_or_build_python_kv_caches(kv_caches); + py::object output = py_model_.attr("write_context_kv")(target_hidden, + positions, + device_cache_slots, + python_kv_caches, + layer_synchronizer); + if (output.is_none()) { + return ModelOutput(); + } + return ModelOutput(output.cast()); +} + +torch::Tensor PyCausalLM::dspark_markov_bias( + const torch::Tensor& previous_token_ids) { + torch::NoGradGuard no_grad; + py::gil_scoped_acquire gil; + return py_model_.attr("dspark_markov_bias")(previous_token_ids) + .cast(); +} + +torch::Tensor PyCausalLM::dspark_confidence_probs( + const torch::Tensor& hidden_all, + const torch::Tensor& prev_matrix) { + torch::NoGradGuard no_grad; + py::gil_scoped_acquire gil; + py::object previous = prev_matrix.defined() + ? py::object(py::cast(prev_matrix)) + : py::object(py::none()); + return py_model_.attr("dspark_confidence_probs")(hidden_all, previous) + .cast(); +} + +bool PyCausalLM::has_dspark_confidence_head() const { + py::gil_scoped_acquire gil; + return py_model_.attr("has_dspark_confidence_head")().cast(); +} + bool PyCausalLM::share_weights_from(CausalLM& source) { auto* source_model = dynamic_cast(&source); if (source_model == nullptr) { diff --git a/xllm/models/llm/py_causal_lm.h b/xllm/models/llm/py_causal_lm.h index 87c6d23e47..27f838a935 100644 --- a/xllm/models/llm/py_causal_lm.h +++ b/xllm/models/llm/py_causal_lm.h @@ -68,6 +68,21 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { torch::Tensor logits(const torch::Tensor& hidden_states, const torch::Tensor& seleted_idxes) override; + ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) override; + + torch::Tensor dspark_markov_bias( + const torch::Tensor& previous_token_ids) override; + + torch::Tensor dspark_confidence_probs( + const torch::Tensor& hidden_all, + const torch::Tensor& prev_matrix) override; + + bool has_dspark_confidence_head() const override; + void load_model(std::unique_ptr loader) override; torch::Device device() const override { return device_; } @@ -83,6 +98,8 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { private: pybind11::dict build_config_dict(const ParallelArgs& parallel_args) const; + const pybind11::object& get_or_build_python_kv_caches( + std::vector& kv_caches); ModelArgs model_args_; torch::TensorOptions options_; @@ -103,6 +120,8 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { pybind11::object py_model_; pybind11::object config_dict_; + pybind11::object python_kv_caches_; + int64_t python_kv_cache_layer_count_ = 0; }; } // namespace xllm diff --git a/xllm/models/llm/qwen3.h b/xllm/models/llm/qwen3.h index 13079f6dea..a4f0b8ed0f 100644 --- a/xllm/models/llm/qwen3.h +++ b/xllm/models/llm/qwen3.h @@ -278,6 +278,11 @@ REGISTER_MODEL_ARGS(qwen3, [&] { LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); LOAD_ARG_OR(max_window_layers, "max_window_layers", 28); + LOAD_ARG_OR(markov_rank, "markov_rank", 0); + LOAD_ARG_OR(enable_confidence_head, "enable_confidence_head", false); + LOAD_ARG_OR( + confidence_head_with_markov, "confidence_head_with_markov", false); + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { return args->hidden_size() / args->n_heads(); }); diff --git a/xllm/python/attention/npu_paged_attention.py b/xllm/python/attention/npu_paged_attention.py index 50925ec0e2..aec83e5b2e 100644 --- a/xllm/python/attention/npu_paged_attention.py +++ b/xllm/python/attention/npu_paged_attention.py @@ -388,6 +388,8 @@ def execute( # prefix. cp_context = get_forward_context().cp_context if cp_context is not None: + if not layer.causal: + raise NotImplementedError("non-causal draft attention does not support context parallelism") return self._prefill_cp(q_3d, k_3d, v_3d, metadata, cp_context, k_cache, v_cache) # Write KV to paged cache (kernel expects [T, kv_heads, head_dim]). @@ -396,7 +398,16 @@ def execute( if metadata.is_prefill or metadata.is_chunked_prefill: if self._use_expanded_decode: return self._decode(q_3d, k_cache, v_cache, metadata, num_tokens) - return self._prefill(q_3d, k_3d, v_3d, k_cache, v_cache, metadata, num_tokens) + return self._prefill( + q_3d, + k_3d, + v_3d, + k_cache, + v_cache, + metadata, + num_tokens, + layer.causal, + ) return self._decode(q_3d, k_cache, v_cache, metadata, num_tokens) def execute_mla( @@ -741,8 +752,11 @@ def _prefill( v_cache: torch.Tensor, metadata: AttentionMetadata, num_tokens: int, + causal: bool, ) -> torch.Tensor: actual_seq = self._cumulative_seq_lens(metadata, num_tokens) + atten_mask = self._causal_mask if causal else None + sparse_mode = _SPARSE_MODE_RIGHT_DOWN_CAUSAL if causal else _SPARSE_MODE_NONE # Prefix-cache hit (or chunked prefill with prior context): part of the # KV already lives in the paged cache, so this forward only carries the @@ -759,7 +773,7 @@ def _prefill( k_flat, v_flat, pse_shift=None, - atten_mask=self._causal_mask, + atten_mask=atten_mask, block_table=self._block_table_i32, actual_seq_lengths=actual_seq, actual_seq_lengths_kv=self._actual_seq_kv, @@ -768,7 +782,7 @@ def _prefill( input_layout="TND", num_key_value_heads=self.num_kv_heads, block_size=block_size, - sparse_mode=_SPARSE_MODE_RIGHT_DOWN_CAUSAL, + sparse_mode=sparse_mode, softmax_lse_flag=False, ) return output.reshape(num_tokens, self.num_heads * self.head_dim) @@ -778,14 +792,14 @@ def _prefill( k_3d, v_3d, pse_shift=None, - atten_mask=self._causal_mask, + atten_mask=atten_mask, actual_seq_lengths=actual_seq, actual_seq_lengths_kv=actual_seq, num_heads=self.num_heads, scale=self.scale, input_layout="TND", num_key_value_heads=self.num_kv_heads, - sparse_mode=_SPARSE_MODE_RIGHT_DOWN_CAUSAL, + sparse_mode=sparse_mode, softmax_lse_flag=False, ) return output.reshape(num_tokens, self.num_heads * self.head_dim) diff --git a/xllm/python/kernels_cuda/__init__.py b/xllm/python/kernels_cuda/__init__.py index 36cbeca543..592c57fccb 100644 --- a/xllm/python/kernels_cuda/__init__.py +++ b/xllm/python/kernels_cuda/__init__.py @@ -48,6 +48,7 @@ chunk_gated_delta_rule, fused_gdn_prefill_post_conv, fused_recurrent_gated_delta_rule_packed_decode, + gdn_prefill_prepare, resolve_gdn_prefill_backend, ) from .linear import prepare_row_parallel_weight @@ -115,6 +116,7 @@ "causal_conv1d_prefill", "causal_conv1d_decode", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_cuda/gated_delta_net.py b/xllm/python/kernels_cuda/gated_delta_net.py index b7409c3bb7..377bc87c7d 100644 --- a/xllm/python/kernels_cuda/gated_delta_net.py +++ b/xllm/python/kernels_cuda/gated_delta_net.py @@ -199,9 +199,58 @@ def _chunk_gated_delta_rule_fake( return torch.empty_like(v), torch.empty_like(initial_state) +def gdn_prefill_prepare( + mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. + + Encapsulates the CUDA-optimal fusion strategy: causal_conv1d_prefill + produces packed convolved output, then fused_gdn_prefill_post_conv + splits into Q/K/V with l2norm and computes gating in one kernel. + + Returns: + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. + """ + from .causal_conv1d import causal_conv1d_prefill + + convolved = causal_conv1d_prefill( + mixed_qkv, + weight, + conv_state, + state_indices, + has_initial_state, + cu_seqlens, + ) + q, k, v, g, beta = fused_gdn_prefill_post_conv( + convolved, + a, + b, + a_log, + dt_bias, + num_key_heads, + key_head_dim, + value_head_dim, + ) + return q, k, v, g, beta + + __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_npu/__init__.py b/xllm/python/kernels_npu/__init__.py index 5626492d7d..8bafc21d3f 100644 --- a/xllm/python/kernels_npu/__init__.py +++ b/xllm/python/kernels_npu/__init__.py @@ -34,11 +34,12 @@ "update_decode_graph_metadata", "vision_fusion_attention", ), - "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_prefill"), + "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_qkv_prefill"), "gated_delta_net": ( "chunk_gated_delta_rule", - "fused_gdn_prefill_post_conv", + "fused_gdn_gating", "fused_recurrent_gated_delta_rule_packed_decode", + "gdn_prefill_prepare", "resolve_gdn_prefill_backend", ), "linear": ("prepare_quant_weight", "prepare_row_parallel_weight"), @@ -133,10 +134,9 @@ "scatter_nd_update", "sparse_flash_attention", "sparse_flash_attention_out", - "causal_conv1d_prefill", "causal_conv1d_decode", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/_custom_op.py b/xllm/python/kernels_npu/_custom_op.py index 2957a5cc85..60e9a0dd64 100644 --- a/xllm/python/kernels_npu/_custom_op.py +++ b/xllm/python/kernels_npu/_custom_op.py @@ -66,6 +66,77 @@ def _rms_norm_fake( return torch.empty_like(input) +def _rms_norm_gated_fake( + input: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + del gate, weight, eps + return torch.empty_like(input) + + +def _l2_norm_fake( + input: torch.Tensor, + eps: float, +) -> torch.Tensor: + del eps + return torch.empty_like(input) + + +def _chunk_gated_delta_rule_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del q, k, g, beta, cu_seqlens + num_seqs = initial_state.shape[0] + return torch.empty_like(v), torch.empty_like(initial_state) + + +def _causal_conv1d_qkv_prefill_fake( + x: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + query_start_loc: torch.Tensor, + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del weight, conv_state, state_indices, has_initial_state, query_start_loc + num_tokens = x.shape[0] + opts = x.options().dtype(torch.bfloat16) + q = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + k = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + v = torch.empty(1, num_tokens, num_v_heads, head_v_dim, **opts) + return q, k, v + + +def _fused_sigmoid_gating_delta_rule_decode_fake( + a_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + ssm_state: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float, +) -> torch.Tensor: + del a_log, a, dt_bias, k, b, ssm_state, state_indices, cu_seqlens, scale + # Output shape matches v: [num_tokens, num_value_heads, value_dim] + return torch.empty_like(v) + + def _fused_add_rms_norm_fake( input: torch.Tensor, residual: torch.Tensor, @@ -523,6 +594,14 @@ def _sparse_flash_attention_out_fake( register_fake("xllm_ops::rms_norm", _rms_norm_fake) +register_fake("xllm_ops::rms_norm_gated", _rms_norm_gated_fake) +register_fake("xllm_ops::l2_norm", _l2_norm_fake) +register_fake("xllm_ops::chunk_gated_delta_rule", _chunk_gated_delta_rule_fake) +register_fake("xllm_ops::causal_conv1d_qkv_prefill", _causal_conv1d_qkv_prefill_fake) +register_fake( + "xllm_ops::fused_sigmoid_gating_delta_rule_decode", + _fused_sigmoid_gating_delta_rule_decode_fake, +) register_fake("xllm_ops::fused_add_rms_norm", _fused_add_rms_norm_fake) register_fake("xllm_ops::silu_and_mul", _silu_and_mul_fake) register_fake("xllm_ops::reshape_paged_cache", _reshape_paged_cache_fake) diff --git a/xllm/python/kernels_npu/causal_conv1d.py b/xllm/python/kernels_npu/causal_conv1d.py index 5f2fbe3a19..d506260d7d 100644 --- a/xllm/python/kernels_npu/causal_conv1d.py +++ b/xllm/python/kernels_npu/causal_conv1d.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU causal-convolution kernels. +"""NPU causal-convolution kernels (PyTorch small-op implementation). -Neither has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/causal_conv1d.py`` and the -Triton launcher it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton reference in +``kernels_cuda/triton/causal_conv1d.py`` using only standard PyTorch +operations. Performance is not optimized; correctness and precision +alignment are the goals. """ from __future__ import annotations @@ -24,32 +25,35 @@ import torch -def causal_conv1d_prefill( +def causal_conv1d_qkv_prefill( value: torch.Tensor, weight: torch.Tensor, conv_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, query_start_loc: torch.Tensor, -) -> torch.Tensor: - """Convolve a variable-length batch and update the convolution states. - - Args: - value: Packed activations of shape ``[num_tokens, channels]``. - weight: Depthwise kernel of shape ``[channels, kernel_size]``. - conv_state: Per-sequence convolution state, updated in place. - state_indices: State slot of every sequence. - has_initial_state: Whether a sequence continues an earlier state. - query_start_loc: Start offset of every sequence in ``value``. + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split into Q/K/V for prefill. Returns: - Convolved activations with the shape and dtype of ``value``. + (q, k, v) with shapes [1, T, num_qk_heads, head_k_dim], + [1, T, num_qk_heads, head_k_dim], [1, T, num_v_heads, head_v_dim]. """ - del value, weight, conv_state, state_indices, has_initial_state - del query_start_loc - raise NotImplementedError( - "causal_conv1d_prefill has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + return torch.ops.xllm_ops.causal_conv1d_qkv_prefill( + value, + weight, + conv_state, + state_indices, + has_initial_state.to(torch.int64), + query_start_loc, + num_qk_heads, + num_v_heads, + head_k_dim, + head_v_dim, ) @@ -70,11 +74,20 @@ def causal_conv1d_decode( Returns: Convolved activations with the shape and dtype of ``value``. """ - del value, weight, conv_state, state_indices - raise NotImplementedError( - "causal_conv1d_decode has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + from .tilelang.causal_conv1d_decode import causal_conv1d_decode as _tl_decode + + # TileLang expects conv_state as [slots, dim, state_len] (PyTorch convention) + # Our cache is [slots, state_len, dim], so transpose before calling + conv_state_pt = conv_state.transpose(1, 2).contiguous() + result = _tl_decode( + x=value, + conv_state=conv_state_pt, + weight=weight, + conv_state_indices=state_indices, ) + # TileLang writes back to conv_state_pt, copy back to original layout + conv_state.copy_(conv_state_pt.transpose(1, 2)) + return result -__all__ = ["causal_conv1d_prefill", "causal_conv1d_decode"] +__all__ = ["causal_conv1d_qkv_prefill", "causal_conv1d_decode"] diff --git a/xllm/python/kernels_npu/gated_delta_net.py b/xllm/python/kernels_npu/gated_delta_net.py index 3fa3577a38..f4f19d74df 100644 --- a/xllm/python/kernels_npu/gated_delta_net.py +++ b/xllm/python/kernels_npu/gated_delta_net.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU gated-delta-network kernels. +"""NPU gated-delta-network kernels (PyTorch small-op implementation). -None has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/gated_delta_net.py`` and the -Triton launchers it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton references in +``kernels_cuda/triton/gdn_prefill.py`` and ``kernels_cuda/triton/gated_delta_net.py`` +using only standard PyTorch operations. Performance is not optimized; +correctness and precision alignment are the goals. """ from __future__ import annotations @@ -25,64 +26,97 @@ import torch -GdnPrefillBackend = Literal["flashinfer", "triton"] +GdnPrefillBackend = Literal["pytorch_naive"] def resolve_gdn_prefill_backend( capability: tuple[int, int] | None = None, ) -> GdnPrefillBackend: - """Select the prefill backend of the active device. + """Select the prefill backend for NPU. Args: - capability: Device capability to resolve for; ``None`` reads it from - the current device. + capability: Ignored on NPU. Returns: - The name to pass as ``backend`` to :func:`chunk_gated_delta_rule`. + The backend name to pass to :func:`chunk_gated_delta_rule`. """ del capability - raise NotImplementedError( - "resolve_gdn_prefill_backend has no NPU implementation; gated delta networks are not supported on NPU yet" + return "pytorch_naive" + + +def fused_gdn_gating( + a_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute decay gate g and beta from raw projections via TileLang kernel.""" + from .tilelang.fused_gdn_gating import fused_gdn_gating_kernel_jit + + num_batches, num_heads = a.shape + kernel = fused_gdn_gating_kernel_jit( + num_batches=num_batches, + compile_max_batch=num_batches, + num_heads=num_heads, ) + g_out = torch.empty(1, num_batches, num_heads, dtype=torch.float32, device=a.device) + beta_out = torch.empty(1, num_batches, num_heads, dtype=a.dtype, device=a.device) + kernel( + a_log.to(torch.float32).contiguous(), + a.contiguous(), + b.contiguous(), + dt_bias.to(torch.float32).contiguous(), + g_out.squeeze(0), + beta_out.squeeze(0), + num_batches, + 1.0, # softplus_beta + 20.0, # softplus_threshold + ) + return g_out, beta_out + -def fused_gdn_prefill_post_conv( +def gdn_prefill_prepare( mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, a: torch.Tensor, b: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor, num_key_heads: int, + num_value_heads: int, key_head_dim: int, value_head_dim: int, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Split the post-convolution projection and build the recurrence gates. +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. - Args: - mixed_qkv: Packed projection of shape ``[num_tokens, qkv_size]``. - a: Gate projection of shape ``[num_tokens, num_value_heads]``. - b: Beta projection of shape ``[num_tokens, num_value_heads]``. - a_log: Per-head log decay of shape ``[num_value_heads]``. - dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. - num_key_heads: Key heads on this rank. - key_head_dim: Size of one key head. - value_head_dim: Size of one value head. + Encapsulates the NPU-optimal fusion strategy: causal_conv1d_qkv does + conv + split + l2norm in one kernel, then fused_gdn_gating computes + decay and beta independently. Returns: - Query, key, value, the decay gate and beta. + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. """ - del mixed_qkv, a, b, a_log, dt_bias - del num_key_heads, key_head_dim, value_head_dim - raise NotImplementedError( - "fused_gdn_prefill_post_conv has no NPU kernel; see " - "kernels_cuda/triton/gdn_prefill.py for the reference implementation" + from .causal_conv1d import causal_conv1d_qkv_prefill + + q, k, v = causal_conv1d_qkv_prefill( + mixed_qkv, + weight, + conv_state, + state_indices, + has_initial_state, + cu_seqlens, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, ) + g, beta = fused_gdn_gating(a_log, a, b, dt_bias) + return q.squeeze(0), k.squeeze(0), v.squeeze(0), g.squeeze(0), beta.squeeze(0) def fused_recurrent_gated_delta_rule_packed_decode( @@ -104,17 +138,52 @@ def fused_recurrent_gated_delta_rule_packed_decode( a_log: Per-head log decay of shape ``[num_value_heads]``. dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. initial_state: Recurrent state pool, updated in place. + Shape ``[num_slots, num_value_heads, value_dim, key_dim]``. state_indices: State slot of every sequence. scale: Query scale. Returns: Output of shape ``[batch_size, 1, num_value_heads, value_head_dim]``. """ - del mixed_qkv, a, b, a_log, dt_bias, initial_state, state_indices, scale - raise NotImplementedError( - "fused_recurrent_gated_delta_rule_packed_decode has no NPU kernel; see " - "kernels_cuda/triton/gated_delta_net.py for the reference implementation" + batch = mixed_qkv.shape[0] + num_value_heads, value_dim, key_dim = initial_state.shape[-3:] + qkv_dim = mixed_qkv.shape[1] + query_key_dim = qkv_dim - num_value_heads * value_dim + query_dim = query_key_dim // 2 + num_key_heads = query_dim // key_dim + + # Split mixed_qkv + q_flat = mixed_qkv[:, :query_dim] + k_flat = mixed_qkv[:, query_dim : 2 * query_dim] + v_flat = mixed_qkv[:, 2 * query_dim :] + + q = q_flat.view(batch, num_key_heads, key_dim) + k = k_flat.view(batch, num_key_heads, key_dim) + v = v_flat.view(batch, num_value_heads, value_dim) + + # Kernel expects [batch, seq_len, heads, dim] — add seq dim for decode + q = q.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + k = k.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + v = v.unsqueeze(1) # [batch, 1, num_value_heads, value_dim] + + # Kernel does l2norm, gating, GQA expansion, and recurrence internally. + # cu_seqlens for decode: each seq has 1 token. + cu_seqlens = torch.arange(batch + 1, dtype=torch.int32, device=mixed_qkv.device) + + output = torch.ops.xllm_ops.fused_sigmoid_gating_delta_rule_decode( + a_log, + a.unsqueeze(1), + dt_bias, + q.contiguous(), + k.contiguous(), + v.contiguous(), + b.unsqueeze(1), + initial_state, + state_indices, + cu_seqlens, + scale, ) + return output.unsqueeze(1) def chunk_gated_delta_rule( @@ -136,22 +205,34 @@ def chunk_gated_delta_rule( g: Decay gate of shape ``[num_tokens, num_value_heads]``. beta: Beta with the shape of ``g``. initial_state: Recurrent state each sequence starts from. + Shape ``[batch, num_value_heads, value_dim, key_dim]``. cu_seqlens: Cumulative sequence lengths. - backend: Name returned by :func:`resolve_gdn_prefill_backend`. + backend: Ignored on NPU. Returns: The output with the shape of ``v`` and the final recurrent state. """ - del q, k, v, g, beta, initial_state, cu_seqlens, backend - raise NotImplementedError( - "chunk_gated_delta_rule has no NPU kernel; see kernels_cuda/triton/fla/ for the reference implementation" + del backend + # npu_mega_chunk_gdn expects [B, T, H, D] layout with B=1 for packed input + # Cast g and beta to match C++ layer behavior (bf16 round-trip) + g_input = g.to(v.dtype) + beta_input = beta.to(v.dtype) + output, final_state = torch.ops.xllm_ops.chunk_gated_delta_rule( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + g_input.unsqueeze(0), + beta_input.unsqueeze(0), + initial_state, + cu_seqlens, ) + return output.squeeze(0), final_state __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/normalization.py b/xllm/python/kernels_npu/normalization.py index 19164de26a..1e967c3dd8 100644 --- a/xllm/python/kernels_npu/normalization.py +++ b/xllm/python/kernels_npu/normalization.py @@ -51,10 +51,7 @@ def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: Returns: A tensor with the shape and dtype of ``value``. """ - del value, eps - raise NotImplementedError( - "l2_norm has no NPU kernel; see kernels_cuda/triton/l2_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.l2_norm(value, eps) def rms_norm_gated( @@ -63,21 +60,18 @@ def rms_norm_gated( weight: torch.Tensor, eps: float = 1e-6, ) -> torch.Tensor: - """Apply RMSNorm to ``value`` and gate the result with ``gate``. + """Apply RMSNorm to ``value`` and gate the result with ``silu(gate)``. Args: value: Tensor to normalize. - gate: Gate applied after normalization, same shape as ``value``. + gate: Gate applied after normalization (SiLU is applied internally). weight: RMSNorm weight over the last dimension. eps: RMSNorm epsilon. Returns: A tensor with the shape and dtype of ``value``. """ - del value, gate, weight, eps - raise NotImplementedError( - "rms_norm_gated has no NPU kernel; see kernels_cuda/triton/rms_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.rms_norm_gated(value, gate, weight, eps) __all__ = [ diff --git a/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py b/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py index aa3815ef7c..8a43ebae94 100644 --- a/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py +++ b/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py @@ -17,8 +17,6 @@ tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True, } -_decode_kernel_cache = {} - def build_causal_conv1d_decode_kernel( width: int, @@ -168,31 +166,6 @@ def _build_decode_kernel_jit( ) -def get_decode_kernel( - width: int, - dim: int, - dtype_str: str = "bfloat16", - has_silu: bool = True, -) -> torch.nn.Module: - dim_chunks = (dim + DIM_PER_CORE - 1) // DIM_PER_CORE - cache_key = ( - width, - dim_chunks, - DIM_PER_CORE, - dtype_str, - has_silu, - ) - if cache_key not in _decode_kernel_cache: - _decode_kernel_cache[cache_key] = _build_decode_kernel_jit( - width, - dim_chunks, - DIM_PER_CORE, - dtype_str, - has_silu, - ) - return _decode_kernel_cache[cache_key] - - def causal_conv1d_decode( x: torch.Tensor, conv_state: torch.Tensor, @@ -268,7 +241,8 @@ def causal_conv1d_decode( initial_state_mode = torch.ones(batch, dtype=torch.int32, device=conv_state.device) - kernel = get_decode_kernel(width, dim, "bfloat16", has_silu) + dim_chunks = (dim + DIM_PER_CORE - 1) // DIM_PER_CORE + kernel = _build_decode_kernel_jit(width, dim_chunks, DIM_PER_CORE, "bfloat16", has_silu) output = kernel( x_kernel, weight_t, @@ -288,3 +262,31 @@ def causal_conv1d_decode( output = output.to(torch.float16) return output + + +if __name__ == "__main__": + import torch + import torch_npu + + batch = 1 + dim = 8192 + kernel_width = 4 + state_len = kernel_width - 1 # 3 + slots = 202 + + device = "npu:0" + tilelang.disable_cache() + tilelang.cache.clear_cache() + + x = torch.randn(batch, dim, dtype=torch.bfloat16, device=device) + weight = torch.randn(dim, kernel_width, dtype=torch.bfloat16, device=device) + # conv_state in PyTorch convention: [slots, dim, state_len] + conv_state = torch.zeros(slots, dim, state_len, dtype=torch.bfloat16, device=device) + state_indices = torch.tensor([1], dtype=torch.int32, device=device) + + print(f"x={x.shape}, weight={weight.shape}, conv_state={conv_state.shape}, state_indices={state_indices}") + print("Calling causal_conv1d_decode...") + result = causal_conv1d_decode(x, conv_state, weight, conv_state_indices=state_indices) + print(f"result={result.shape}, dtype={result.dtype}") + print(f"result first 5: {result[0, :5].tolist()}") + print("SUCCESS") diff --git a/xllm/python/layers/attention.py b/xllm/python/layers/attention.py index 947a3368eb..f1336fc556 100644 --- a/xllm/python/layers/attention.py +++ b/xllm/python/layers/attention.py @@ -37,6 +37,7 @@ def __init__( scale: float, sliding_window: int, layer_id: int, + causal: bool = True, ) -> None: super().__init__() self.num_heads = num_heads @@ -45,6 +46,7 @@ def __init__( self.scale = scale self.sliding_window = sliding_window self.layer_id = layer_id + self.causal = causal def forward( self, diff --git a/xllm/python/layers/gated_delta_net.py b/xllm/python/layers/gated_delta_net.py index cd081b396e..02667f71c2 100644 --- a/xllm/python/layers/gated_delta_net.py +++ b/xllm/python/layers/gated_delta_net.py @@ -125,15 +125,25 @@ def _conv_prefill( state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, - ) -> torch.Tensor: + a: torch.Tensor, + b: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: self._conv_state_dim_first(conv_state) - return kernels.causal_conv1d_prefill( + return kernels.gdn_prefill_prepare( mixed_qkv, self.conv1d_weight, conv_state, state_indices, has_initial_state, cu_seqlens, + a, + b, + self.A_log, + self.dt_bias, + self.num_k_heads, + self.num_v_heads, + self.key_head_dim, + self.value_head_dim, ) def _conv_decode( @@ -152,17 +162,16 @@ def _conv_decode( def _gdn_prefill( self, - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, ssm_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, ) -> torch.Tensor: - # TODO: Fuse cache gather/zero/scatter and null-row output masking into - # the GDN kernels. The current staging is intentionally kept for this - # correctness PR and should be removed in the next performance PR. non_null_state = state_indices > 0 use_initial_state = non_null_state & has_initial_state cache_indices = state_indices.to(torch.long) @@ -173,16 +182,6 @@ def _gdn_prefill( initial_state, torch.zeros_like(initial_state), ) - q, k, v, g, beta = kernels.fused_gdn_prefill_post_conv( - mixed_qkv=mixed_qkv, - a=a, - b=b, - a_log=self.A_log, - dt_bias=self.dt_bias, - num_key_heads=self.num_k_heads, - key_head_dim=self.key_head_dim, - value_head_dim=self.value_head_dim, - ) output, final_state = kernels.chunk_gated_delta_rule( q, k, @@ -193,11 +192,12 @@ def _gdn_prefill( cu_seqlens, self.gdn_prefill_backend, ) + num_tokens = q.shape[0] sequence_lengths = cu_seqlens.diff().to(dtype=torch.long) token_mask = torch.repeat_interleave( non_null_state, sequence_lengths, - output_size=mixed_qkv.shape[0], + output_size=num_tokens, ) output = torch.where(token_mask[:, None, None], output, 0.0) ssm_state.index_copy_( @@ -256,21 +256,25 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: has_initial_state = has_initial_state.to(device=hidden.device, dtype=torch.bool) if has_initial_state.shape != state_indices.shape: raise ValueError("has_initial_state must match linear_state_indices") - mixed_qkv = self._conv_prefill( + q, k, v, g, beta = self._conv_prefill( mixed_qkv, conv_state, state_indices, has_initial_state, cu_seqlens, + a, + b, ) else: mixed_qkv = self._conv_decode(mixed_qkv, conv_state, state_indices) if is_prefill: output = self._gdn_prefill( - mixed_qkv, - a, - b, + q, + k, + v, + g, + beta, ssm_state, state_indices, has_initial_state, diff --git a/xllm/python/model_executor/executor.py b/xllm/python/model_executor/executor.py index 0d8dd5f6cc..8c2a623637 100644 --- a/xllm/python/model_executor/executor.py +++ b/xllm/python/model_executor/executor.py @@ -25,6 +25,7 @@ ) from xllm.python.layers.attention import Attention from xllm.python.model_executor.forward_context import LayerSynchronizer +from xllm.python.model_executor.runners.base import ModelExecutionOutput from xllm.python.model_executor.runners.eager import EagerRunner from xllm.python.platform import current_platform @@ -212,7 +213,7 @@ def execute( metadata: AttentionMetadata, input_embedding: torch.Tensor | None = None, layer_synchronizer: LayerSynchronizer | None = None, - ) -> torch.Tensor: + ) -> ModelExecutionOutput: if not self._kv_bound: raise RuntimeError("KV caches are not bound") diff --git a/xllm/python/model_executor/forward_context.py b/xllm/python/model_executor/forward_context.py index ebd1f92790..f6b7786f25 100644 --- a/xllm/python/model_executor/forward_context.py +++ b/xllm/python/model_executor/forward_context.py @@ -39,7 +39,7 @@ class LayerSynchronizer(Protocol): forward to finish. """ - def record_event(self, layer_id: int) -> None: ... + def record_event(self, layer_id: int) -> bool: ... @dataclass(frozen=True, slots=True) diff --git a/xllm/python/model_executor/runners/base.py b/xllm/python/model_executor/runners/base.py index 1dd8932b75..431247f75e 100644 --- a/xllm/python/model_executor/runners/base.py +++ b/xllm/python/model_executor/runners/base.py @@ -26,6 +26,8 @@ ) from xllm.python.model_executor.forward_context import LayerSynchronizer +ModelExecutionOutput = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + class BaseRunner(ABC): def __init__( @@ -50,5 +52,5 @@ def execute( metadata: AttentionMetadata, input_embedding: torch.Tensor | None = None, layer_synchronizer: LayerSynchronizer | None = None, - ) -> torch.Tensor: + ) -> ModelExecutionOutput: pass diff --git a/xllm/python/model_executor/runners/eager.py b/xllm/python/model_executor/runners/eager.py index 5540ad5d1a..26128d5da3 100644 --- a/xllm/python/model_executor/runners/eager.py +++ b/xllm/python/model_executor/runners/eager.py @@ -23,7 +23,7 @@ LayerSynchronizer, forward_context, ) -from xllm.python.model_executor.runners.base import BaseRunner +from xllm.python.model_executor.runners.base import BaseRunner, ModelExecutionOutput def _per_seq_lens_from_metadata(metadata: AttentionMetadata) -> list[int] | None: @@ -53,7 +53,7 @@ def execute( metadata: AttentionMetadata, input_embedding: torch.Tensor | None = None, layer_synchronizer: LayerSynchronizer | None = None, - ) -> torch.Tensor: + ) -> ModelExecutionOutput: self.attention_backend.prepare(metadata) cp_context = None diff --git a/xllm/python/model_platform_support.py b/xllm/python/model_platform_support.py index caae712526..b7721ad928 100644 --- a/xllm/python/model_platform_support.py +++ b/xllm/python/model_platform_support.py @@ -16,7 +16,9 @@ MODEL_PLATFORM_SUPPORT: dict[str, dict[str, bool]] = { "qwen3": {"cuda": True, "npu": True}, - "qwen3_5": {"cuda": True, "npu": False}, + "qwen3_5": {"cuda": True, "npu": True}, + "qwen3_dflash": {"cuda": False, "npu": True}, + "qwen3_dspark": {"cuda": False, "npu": True}, "qwen3_vl": {"cuda": False, "npu": True}, "deepseek_v32": {"cuda": False, "npu": True}, "glm5_2": {"cuda": False, "npu": True}, diff --git a/xllm/python/models/aux_hidden_capture.py b/xllm/python/models/aux_hidden_capture.py new file mode 100644 index 0000000000..0b4b29a710 --- /dev/null +++ b/xllm/python/models/aux_hidden_capture.py @@ -0,0 +1,52 @@ +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Intermediate residual-stream capture for Python target models.""" + +from __future__ import annotations + +import torch + + +class AuxHiddenCapture: + """Captures selected layer inputs and concatenates them in config order.""" + + def __init__(self, layers_to_capture: tuple[int, ...]) -> None: + self._layers_to_capture = layers_to_capture + self._capture_set = frozenset(layers_to_capture) + + @property + def enabled(self) -> bool: + return bool(self._layers_to_capture) + + def capture_layer( + self, + layer_id: int, + hidden: torch.Tensor, + residual: torch.Tensor | None, + captured: dict[int, torch.Tensor], + ) -> None: + if layer_id not in self._capture_set: + return + captured[layer_id] = hidden.clone() if residual is None else hidden + residual + + def finalize( + self, + hidden: torch.Tensor, + captured: dict[int, torch.Tensor], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if not self.enabled: + return hidden + aux_hidden = torch.cat([captured[layer_id] for layer_id in self._layers_to_capture], dim=-1) + return hidden, aux_hidden diff --git a/xllm/python/models/dspark.py b/xllm/python/models/dspark.py new file mode 100644 index 0000000000..7910306a2c --- /dev/null +++ b/xllm/python/models/dspark.py @@ -0,0 +1,116 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared DSpark heads and runtime methods for Python model adapters.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from xllm.python.models.base import PyModelBase + + +class DSparkMarkovHead(nn.Module): + def __init__( + self, + vocab_size: int, + draft_vocab_size: int, + markov_rank: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__() + self.markov_w1 = nn.Embedding(vocab_size, markov_rank, dtype=dtype, device=device) + self.markov_w2 = nn.Linear(markov_rank, draft_vocab_size, bias=False, dtype=dtype, device=device) + + def embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w1(token_ids) + + def bias(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_w2(self.embed(token_ids)) + + +class DSparkConfidenceHead(nn.Module): + def __init__( + self, + hidden_size: int, + markov_rank: int, + with_markov: bool, + device: torch.device, + ) -> None: + super().__init__() + self.with_markov = with_markov + input_size = hidden_size + markov_rank if with_markov else hidden_size + self.proj = nn.Linear(input_size, 1, bias=True, dtype=torch.float32, device=device) + + def forward(self, hidden: torch.Tensor, markov_embedding: torch.Tensor | None) -> torch.Tensor: + if self.with_markov: + if markov_embedding is None: + raise ValueError("DSpark confidence head requires Markov embeddings") + hidden = torch.cat((hidden, markov_embedding), dim=-1) + return torch.sigmoid(self.proj(hidden.float())).squeeze(-1).to(torch.float32) + + +class DSparkForCausalLMBase(PyModelBase): + """Owns model-independent DSpark heads and runtime bridge methods.""" + + def __init__( + self, + *, + vocab_size: int, + draft_vocab_size: int, + markov_rank: int, + hidden_size: int, + enable_confidence_head: bool, + confidence_head_with_markov: bool, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__() + self.markov_head = DSparkMarkovHead( + vocab_size, + draft_vocab_size, + markov_rank, + dtype, + device, + ) + self.confidence_head = ( + DSparkConfidenceHead( + hidden_size, + markov_rank, + confidence_head_with_markov, + device, + ) + if enable_confidence_head + else None + ) + + def dspark_markov_bias(self, previous_token_ids: torch.Tensor) -> torch.Tensor: + return self.markov_head.bias(previous_token_ids) + + def dspark_confidence_probs( + self, + hidden_all: torch.Tensor, + prev_matrix: torch.Tensor | None, + ) -> torch.Tensor: + if self.confidence_head is None: + raise RuntimeError("DSpark confidence head is not enabled") + markov_embedding = None + if prev_matrix is not None: + markov_embedding = self.markov_head.embed(prev_matrix) + return self.confidence_head(hidden_all, markov_embedding) + + def has_dspark_confidence_head(self) -> bool: + return self.confidence_head is not None diff --git a/xllm/python/models/glm5_2.py b/xllm/python/models/glm5_2.py index dc7c4bddc2..2cddedc6fe 100644 --- a/xllm/python/models/glm5_2.py +++ b/xllm/python/models/glm5_2.py @@ -646,13 +646,11 @@ def load_weights( self.model.layers[i].mlp.process_weights_after_loading() else: se = p + "mlp.experts." + moe_layer = self.model.layers[i].mlp + moe_layer.allocate_experts_w13_for_loading() w13_param = self.get_parameter(p + "mlp.experts_w13") - w2_param = self.get_parameter(p + "mlp.experts_w2") w13_scale = self.get_buffer(p + "mlp.experts_w13_scale") w13_offset = self.get_buffer(p + "mlp.experts_w13_offset") - w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") - w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") - moe_layer = self.model.layers[i].mlp expert_start = moe_layer.local_expert_start expert_end = moe_layer.local_expert_end shard_world = cfg.moe_tp_size if cfg.ep_size > 1 else cfg.tp_size @@ -665,9 +663,6 @@ def load_weights( uw = loader.load_tensor(se + f"{j}.up_proj.weight") us = loader.load_tensor(se + f"{j}.up_proj.weight_scale") uo = loader.load_tensor(se + f"{j}.up_proj.weight_offset") - dw = loader.load_tensor(se + f"{j}.down_proj.weight") - ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") - do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w13_param.data[local_idx].copy_( torch.cat( [ @@ -695,6 +690,16 @@ def load_weights( dim=0, ).contiguous() ) + + moe_layer.allocate_experts_w2_for_loading() + w2_param = self.get_parameter(p + "mlp.experts_w2") + w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") + w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") + for j in range(expert_start, expert_end): + local_idx = j - expert_start + dw = loader.load_tensor(se + f"{j}.down_proj.weight") + ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") + do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w2_param.data[local_idx].copy_(loader.shard(dw, 1, shard_world, shard_rank).contiguous()) w2_scale.data[local_idx].copy_(ds.contiguous()) w2_offset.data[local_idx].copy_(do.contiguous()) diff --git a/xllm/python/models/qwen3.py b/xllm/python/models/qwen3.py index 5ca9b0207a..36c7191dff 100644 --- a/xllm/python/models/qwen3.py +++ b/xllm/python/models/qwen3.py @@ -43,6 +43,7 @@ get_forward_context, record_layer_event, ) # noqa: F401 +from xllm.python.models.aux_hidden_capture import AuxHiddenCapture from xllm.python.models.base import PyModelBase @@ -65,6 +66,7 @@ class Qwen3Config: tp_rank: int = 0 dp_size: int = 1 dp_rank: int = 0 + layers_to_capture: tuple[int, ...] = () @classmethod def from_dict(cls, d: dict) -> Qwen3Config: @@ -94,6 +96,7 @@ def pick(*keys, default=None): tp_rank=int(pick("tp_rank", default=0)), dp_size=int(pick("dp_size", default=1)), dp_rank=int(pick("dp_rank", default=0)), + layers_to_capture=tuple(int(layer_id) for layer_id in pick("layers_to_capture", default=[])), ) def head_split(self) -> tuple[int, int, int]: @@ -113,7 +116,14 @@ def head_split(self) -> tuple[int, int, int]: class Qwen3Attention(nn.Module): - def __init__(self, cfg: Qwen3Config, layer_id: int, dtype: torch.dtype, device: torch.device) -> None: + def __init__( + self, + cfg: Qwen3Config, + layer_id: int, + dtype: torch.dtype, + device: torch.device, + causal: bool = True, + ) -> None: super().__init__() self.layer_id = layer_id num_heads, num_kv_heads, replicas = cfg.head_split() @@ -149,6 +159,7 @@ def __init__(self, cfg: Qwen3Config, layer_id: int, dtype: torch.dtype, device: scale=self.head_dim**-0.5, sliding_window=cfg.sliding_window, layer_id=layer_id, + causal=causal, ) def forward( @@ -217,11 +228,12 @@ def __init__( layer_id: int, dtype: torch.dtype, device: torch.device, + causal: bool = True, ) -> None: super().__init__() self.layer_id = layer_id self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) - self.self_attn = Qwen3Attention(cfg, layer_id, dtype, device) + self.self_attn = Qwen3Attention(cfg, layer_id, dtype, device, causal=causal) self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) self.mlp = GatedMLP( cfg.hidden_size, @@ -271,13 +283,19 @@ def __init__(self, cfg: Qwen3Config, dtype: torch.dtype, device: torch.device) - ) self.layers = nn.ModuleList([Qwen3DecoderLayer(cfg, i, dtype, device) for i in range(cfg.n_layers)]) self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) + invalid_capture_layers = [layer_id for layer_id in cfg.layers_to_capture if layer_id >= cfg.n_layers] + if invalid_capture_layers: + raise ValueError( + f"layers_to_capture must be smaller than n_layers ({cfg.n_layers}): {invalid_capture_layers}" + ) + self.aux_hidden_capture = AuxHiddenCapture(cfg.layers_to_capture) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, mrope_section: list[int] | None = None, - ) -> torch.Tensor: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: hidden = self.embed_tokens(input_ids) # The fused QK-norm+RoPE kernel requires int64 position ids, but C++ # passes them as int32. Cast once here instead of once per layer. In the @@ -290,10 +308,14 @@ def forward( # active on prefill with cp_size>1; cp_context is None otherwise. cp_context = get_forward_context().cp_context if cp_context is not None: + if self.aux_hidden_capture.enabled: + raise NotImplementedError("aux hidden capture does not support context parallelism") hidden = cp_shard_rows(hidden, cp_context) positions = cp_shard_positions(positions, cp_context).contiguous() residual: torch.Tensor | None = None + captured_hidden: dict[int, torch.Tensor] = {} for i, layer in enumerate(self.layers): + self.aux_hidden_capture.capture_layer(i, hidden, residual, captured_hidden) hidden, residual = layer( hidden, residual, @@ -307,7 +329,7 @@ def forward( hidden, _ = self.norm(hidden, residual) if cp_context is not None: hidden = cp_merge_rows(hidden, cp_context) - return hidden + return self.aux_hidden_capture.finalize(hidden, captured_hidden) class Qwen3ForCausalLM(PyModelBase): diff --git a/xllm/python/models/qwen3_5.py b/xllm/python/models/qwen3_5.py index 0c2c12d1d5..6fd43452de 100644 --- a/xllm/python/models/qwen3_5.py +++ b/xllm/python/models/qwen3_5.py @@ -418,6 +418,13 @@ def __init__(self, cfg: Qwen3_5Config, dtype: torch.dtype, device: torch.device) self.norm = GemmaRMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # TODO: The current tilelang-ascend version has a cache bug that prevents kernels with + # dynamic symbols from being cached, causing the service to crash. This is a temporary + # workaround; we will resubmit once tilelang-ascend fixes the issue. + import tilelang + + tilelang.disable_cache() + tilelang.cache.clear_cache() hidden = self.embed_tokens(input_ids) residual: torch.Tensor | None = None for layer in self.layers: diff --git a/xllm/python/models/qwen3_dflash.py b/xllm/python/models/qwen3_dflash.py new file mode 100644 index 0000000000..93c839a6e3 --- /dev/null +++ b/xllm/python/models/qwen3_dflash.py @@ -0,0 +1,393 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3-style DFlash draft model for the Python NPU executor.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from xllm.python import distributed, kernels +from xllm.python.layers import RMSNorm, RotaryEmbedding +from xllm.python.model_executor.forward_context import ( + LayerSynchronizer, + record_layer_event, +) +from xllm.python.models.base import PyModelBase +from xllm.python.models.qwen3 import Qwen3Config, Qwen3DecoderLayer + + +@dataclass +class DFlashQwen3Config(Qwen3Config): + draft_vocab_size: int = 0 + world_size: int = 1 + + @classmethod + def from_dict(cls, d: dict) -> DFlashQwen3Config: + normalized_config = dict(d) + rope_parameters = d.get("rope_parameters") or {} + if normalized_config.get("rope_theta") is None: + normalized_config["rope_theta"] = rope_parameters.get("rope_theta") + + base = Qwen3Config.from_dict(normalized_config) + draft_vocab_size = int(d.get("draft_vocab_size") or base.vocab_size) + return cls( + **base.__dict__, + draft_vocab_size=draft_vocab_size, + world_size=int(d.get("world_size", base.tp_size * base.dp_size)), + ) + + def validate(self) -> None: + if self.hidden_size <= 0 or self.n_layers <= 0 or self.n_heads <= 0: + raise ValueError("invalid Qwen3-style block-diffusion dimensions") + if min(self.tp_size, self.dp_size) <= 0: + raise ValueError("parallel sizes must be positive") + if self.tp_size * self.dp_size != self.world_size: + raise ValueError("world_size must equal tp_size * dp_size") + if not 0 <= self.dp_rank < self.dp_size: + raise ValueError("dp_rank must be in [0, dp_size)") + if not 0 <= self.tp_rank < self.tp_size: + raise ValueError("tp_rank must be in [0, tp_size)") + if self.hidden_size % self.tp_size: + raise ValueError("hidden_size must be divisible by tp_size") + if self.vocab_size % self.tp_size: + raise ValueError("vocab_size must be divisible by tp_size") + if self.draft_vocab_size != self.vocab_size: + raise ValueError("reduced-vocabulary block-diffusion drafts are not supported") + self.head_split() + + +class DFlashContextProjection(nn.Module): + """Column-parallel projection with checkpoint-defined input width.""" + + def __init__( + self, + out_features: int, + tp_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__() + self.out_features = out_features + self.tp_size = tp_size + self.weight = nn.Parameter(torch.empty(out_features // tp_size, 0, dtype=dtype, device=device)) + + def load_weight(self, weight: torch.Tensor, tp_rank: int) -> None: + if weight.dim() != 2 or weight.size(0) != self.out_features: + raise ValueError("DFlash fc.weight has an invalid shape") + local_out_features = self.out_features // self.tp_size + weight = weight.narrow( + 0, + tp_rank * local_out_features, + local_out_features, + ) + self.weight = nn.Parameter(weight.to(dtype=self.weight.dtype, device=self.weight.device).contiguous()) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if self.weight.size(1) == 0: + raise RuntimeError("block-diffusion context projection weight is not loaded") + if hidden.size(-1) != self.weight.size(1): + raise ValueError("target auxiliary hidden size does not match the draft fc.weight") + output = F.linear(hidden, self.weight) + if self.tp_size > 1: + output = distributed.all_gather( + output, + dim=-1, + world_size=self.tp_size, + ) + return output + + +class DFlashQwen3DecoderLayer(Qwen3DecoderLayer): + def __init__( + self, + cfg: DFlashQwen3Config, + layer_id: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__(cfg, layer_id, dtype, device, causal=False) + + +class DFlashQwen3Model(nn.Module): + def __init__(self, cfg: DFlashQwen3Config, dtype: torch.dtype, device: torch.device) -> None: + super().__init__() + self.cfg = cfg + self.embed_tokens: nn.Module | None = None + self.fc = DFlashContextProjection( + cfg.hidden_size, + cfg.tp_size, + dtype, + device, + ) + self.hidden_norm = RMSNorm( + cfg.hidden_size, + cfg.rms_norm_eps, + dtype=dtype, + device=device, + ) + self.rotary = RotaryEmbedding( + cfg.head_dim, + cfg.max_position_embeddings, + cfg.rope_theta, + dtype=dtype, + device=device, + ) + self.layers = nn.ModuleList( + DFlashQwen3DecoderLayer(cfg, layer_id, dtype, device) for layer_id in range(cfg.n_layers) + ) + self.norm = RMSNorm( + cfg.hidden_size, + cfg.rms_norm_eps, + dtype=dtype, + device=device, + ) + self.register_buffer("_fused_kv_weight", None, persistent=False) + self.register_buffer("_fused_kv_bias", None, persistent=False) + self.register_buffer("_k_norm_weights", None, persistent=False) + + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + if self.embed_tokens is None: + raise RuntimeError("target embedding has not been shared with the draft") + hidden = self.embed_tokens(input_ids) + positions = positions.to(torch.int64).contiguous() + residual: torch.Tensor | None = None + for layer_id, layer in enumerate(self.layers): + hidden, residual = layer( + hidden, + residual, + positions, + self.rotary.cos_sin_cache, + None, + None, + ) + record_layer_event(layer_id) + hidden, _ = self.norm(hidden, residual) + return hidden + + def _build_context_kv_buffers(self) -> None: + kv_weights: list[torch.Tensor] = [] + kv_biases: list[torch.Tensor] = [] + k_norm_weights: list[torch.Tensor] = [] + has_bias = self.layers[0].self_attn.qkv_proj.bias is not None + for layer in self.layers: + attention = layer.self_attn + kv_weights.append(attention.qkv_proj.weight[attention.q_size :]) + if has_bias: + kv_biases.append(attention.qkv_proj.bias[attention.q_size :]) + k_norm_weights.append(attention.k_norm.weight) + self._fused_kv_weight = torch.cat(kv_weights, dim=0).detach().contiguous() + self._fused_kv_bias = torch.cat(kv_biases, dim=0).detach().contiguous() if has_bias else None + self._k_norm_weights = torch.stack(k_norm_weights, dim=0).detach().float().view(self.cfg.n_layers, 1, 1, -1) + + def _normalize_context_keys(self, keys: torch.Tensor) -> torch.Tensor: + if self._k_norm_weights is None: + raise RuntimeError("draft K-norm buffers are not initialized") + keys_fp32 = keys.float() + variance = keys_fp32.pow(2).mean(dim=-1, keepdim=True) + normalized = keys_fp32 * torch.rsqrt(variance + self.cfg.rms_norm_eps) + return (normalized * self._k_norm_weights).to(keys.dtype) + + def _apply_context_rope( + self, + keys: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + num_layers, num_context, num_kv_heads, head_dim = keys.shape + flat_keys = keys.reshape(num_layers * num_context, num_kv_heads, head_dim) + repeated_positions = positions.to(torch.long).repeat(num_layers) + cache = self.rotary.cos_sin_cache.index_select(0, repeated_positions) + cos_half, sin_half = cache.chunk(2, dim=-1) + cos = torch.cat((cos_half, cos_half), dim=-1).unsqueeze(1) + sin = torch.cat((sin_half, sin_half), dim=-1).unsqueeze(1) + first, second = flat_keys.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return (flat_keys * cos + rotated * sin).view_as(keys) + + def write_context_kv( + self, + target_hidden: torch.Tensor, + positions: torch.Tensor, + cache_slots: torch.Tensor, + kv_caches: list[tuple[torch.Tensor | None, ...]], + layer_synchronizer: LayerSynchronizer | None, + ) -> torch.Tensor | None: + if len(kv_caches) != self.cfg.n_layers: + raise ValueError("draft KV cache layer count mismatch") + if cache_slots.numel() != target_hidden.size(0): + raise ValueError("draft context cache slot count mismatch") + if self._fused_kv_weight is None or self._k_norm_weights is None: + raise RuntimeError("draft context KV buffers are not initialized") + + projected_hidden = self.hidden_norm(self.fc(target_hidden)) + all_kv = F.linear( + projected_hidden, + self._fused_kv_weight, + self._fused_kv_bias, + ) + num_context = projected_hidden.size(0) + num_kv_heads = self.layers[0].self_attn.num_kv_heads + all_kv = all_kv.view( + num_context, + self.cfg.n_layers, + 2, + num_kv_heads, + self.cfg.head_dim, + ).permute(2, 1, 0, 3, 4) + all_keys = self._apply_context_rope( + self._normalize_context_keys(all_kv[0]), + positions, + ) + all_values = all_kv[1].contiguous() + + for layer_id, cache in enumerate(kv_caches): + key_cache, value_cache = cache[0], cache[1] + if key_cache is None or value_cache is None: + raise RuntimeError(f"draft KV cache is missing for layer {layer_id}") + kernels.reshape_paged_cache( + cache_slots, + all_keys[layer_id].contiguous(), + all_values[layer_id], + key_cache, + value_cache, + ) + if layer_synchronizer is not None: + if layer_synchronizer.record_event(layer_id) is False: + return None + return projected_hidden + + def load_weights(self, state_dicts: list, tp_rank: int, tp_size: int) -> None: + cfg = self.cfg + kv_replicas = tp_size // cfg.n_kv_heads if cfg.n_kv_heads < tp_size else 1 + kv_rank = tp_rank // kv_replicas if kv_replicas > 1 else tp_rank + kv_world = tp_size // kv_replicas if kv_replicas > 1 else tp_size + + def find(name: str): + candidates = (name, f"model.{name}") + for candidate in candidates: + for state_dict in state_dicts: + if state_dict.has(candidate): + return state_dict, candidate + return None + + def load_tensor(name: str) -> torch.Tensor: + found = find(name) + if found is None: + raise KeyError(f"checkpoint tensor not found: {name}") + state_dict, key = found + return state_dict.get_tensor(key) + + def shard(name: str, dim: int, kv: bool = False) -> torch.Tensor: + tensor = load_tensor(name) + rank = kv_rank if kv else tp_rank + world = kv_world if kv else tp_size + if world <= 1: + return tensor + chunk_size = tensor.size(dim) // world + return tensor.narrow(dim, rank * chunk_size, chunk_size).contiguous() + + def copy_in(param_name: str, tensor: torch.Tensor) -> None: + param = self.get_parameter(param_name) + param.data.copy_(tensor.to(dtype=param.dtype, device=param.device)) + + self.fc.load_weight(load_tensor("fc.weight"), tp_rank) + copy_in("hidden_norm.weight", load_tensor("hidden_norm.weight")) + + for layer_id in range(cfg.n_layers): + prefix = f"layers.{layer_id}." + target = f"layers.{layer_id}." + for norm_name in ( + "input_layernorm.weight", + "post_attention_layernorm.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ): + copy_in(target + norm_name, load_tensor(prefix + norm_name)) + + q_weight = shard(prefix + "self_attn.q_proj.weight", dim=0) + k_weight = shard(prefix + "self_attn.k_proj.weight", dim=0, kv=True) + v_weight = shard(prefix + "self_attn.v_proj.weight", dim=0, kv=True) + copy_in( + target + "self_attn.qkv_proj.weight", + torch.cat((q_weight, k_weight, v_weight)), + ) + copy_in( + target + "self_attn.o_proj.weight", + shard(prefix + "self_attn.o_proj.weight", dim=1), + ) + + if cfg.attention_bias: + q_bias = shard(prefix + "self_attn.q_proj.bias", dim=0) + k_bias = shard(prefix + "self_attn.k_proj.bias", dim=0, kv=True) + v_bias = shard(prefix + "self_attn.v_proj.bias", dim=0, kv=True) + copy_in( + target + "self_attn.qkv_proj.bias", + torch.cat((q_bias, k_bias, v_bias)), + ) + copy_in( + target + "self_attn.o_proj.bias", + load_tensor(prefix + "self_attn.o_proj.bias"), + ) + + gate_weight = shard(prefix + "mlp.gate_proj.weight", dim=0) + up_weight = shard(prefix + "mlp.up_proj.weight", dim=0) + copy_in( + target + "mlp.gate_up_proj.weight", + torch.cat((gate_weight, up_weight)), + ) + copy_in( + target + "mlp.down_proj.weight", + shard(prefix + "mlp.down_proj.weight", dim=1), + ) + + layer = self.layers[layer_id] + layer.self_attn.o_proj.process_weights_after_loading() + layer.mlp.down_proj.process_weights_after_loading() + + copy_in("norm.weight", load_tensor("norm.weight")) + self._build_context_kv_buffers() + + +class DFlashQwen3ForCausalLM(PyModelBase): + def __init__(self, config: dict) -> None: + super().__init__() + self.cfg = DFlashQwen3Config.from_dict(config) + self.cfg.validate() + self.dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype")) + self.device = torch.device(config.get("device", "npu")) + self.model = DFlashQwen3Model(self.cfg, self.dtype, self.device) + self.lm_head: nn.Module | None = None + + def load_weights(self, state_dicts: list, tp_rank: int, tp_size: int) -> None: + self.model.load_weights(state_dicts, tp_rank, tp_size) + + def write_context_kv( + self, + target_hidden: torch.Tensor, + positions: torch.Tensor, + cache_slots: torch.Tensor, + kv_caches: list[tuple[torch.Tensor | None, ...]], + layer_synchronizer: LayerSynchronizer | None, + ) -> torch.Tensor | None: + return self.model.write_context_kv( + target_hidden, + positions, + cache_slots, + kv_caches, + layer_synchronizer, + ) diff --git a/xllm/python/models/qwen3_dspark.py b/xllm/python/models/qwen3_dspark.py new file mode 100644 index 0000000000..b7c7904e83 --- /dev/null +++ b/xllm/python/models/qwen3_dspark.py @@ -0,0 +1,127 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3 DSpark draft model for the Python NPU executor.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from xllm.python.model_executor.forward_context import LayerSynchronizer +from xllm.python.models.dspark import DSparkForCausalLMBase +from xllm.python.models.qwen3_dflash import ( + DFlashQwen3Config, + DFlashQwen3Model, +) + + +@dataclass +class Qwen3DSparkConfig(DFlashQwen3Config): + markov_rank: int = 0 + enable_confidence_head: bool = False + confidence_head_with_markov: bool = False + + @classmethod + def from_dict(cls, d: dict) -> Qwen3DSparkConfig: + base = DFlashQwen3Config.from_dict(d) + return cls( + **base.__dict__, + markov_rank=int(d.get("markov_rank", 0)), + enable_confidence_head=bool(d.get("enable_confidence_head", False)), + confidence_head_with_markov=bool(d.get("confidence_head_with_markov", False)), + ) + + def validate(self) -> None: + super().validate() + if self.markov_rank <= 0: + raise ValueError("Qwen3 DSpark requires markov_rank > 0") + + +class Qwen3DSparkModel(DFlashQwen3Model): + """DFlash draft backbone used with DSpark-specific output heads.""" + + +class Qwen3DSparkForCausalLM(DSparkForCausalLMBase): + def __init__(self, config: dict) -> None: + cfg = Qwen3DSparkConfig.from_dict(config) + cfg.validate() + dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype")) + device = torch.device(config.get("device", "npu")) + super().__init__( + vocab_size=cfg.vocab_size, + draft_vocab_size=cfg.draft_vocab_size, + markov_rank=cfg.markov_rank, + hidden_size=cfg.hidden_size, + enable_confidence_head=cfg.enable_confidence_head, + confidence_head_with_markov=cfg.confidence_head_with_markov, + dtype=dtype, + device=device, + ) + self.cfg = cfg + self.dtype = dtype + self.device = device + self.model = Qwen3DSparkModel(cfg, dtype, device) + self.lm_head: nn.Module | None = None + + def load_weights(self, state_dicts: list, tp_rank: int, tp_size: int) -> None: + self.model.load_weights(state_dicts, tp_rank, tp_size) + + def find(name: str): + candidates = (name, f"model.{name}") + for candidate in candidates: + for state_dict in state_dicts: + if state_dict.has(candidate): + return state_dict, candidate + return None + + def load_tensor(name: str) -> torch.Tensor: + found = find(name) + if found is None: + raise KeyError(f"checkpoint tensor not found: {name}") + state_dict, key = found + return state_dict.get_tensor(key) + + def copy_in(param_name: str, tensor: torch.Tensor) -> None: + param = self.get_parameter(param_name) + param.data.copy_(tensor.to(dtype=param.dtype, device=param.device)) + + markov_w1 = load_tensor("markov_head.markov_w1.weight") + markov_w2 = load_tensor("markov_head.markov_w2.weight") + copy_in("markov_head.markov_w1.weight", markov_w1) + copy_in("markov_head.markov_w2.weight", markov_w2) + + if self.confidence_head is not None: + confidence_weight = load_tensor("confidence_head.proj.weight") + confidence_bias = load_tensor("confidence_head.proj.bias") + copy_in("confidence_head.proj.weight", confidence_weight) + copy_in("confidence_head.proj.bias", confidence_bias) + + def write_context_kv( + self, + target_hidden: torch.Tensor, + positions: torch.Tensor, + cache_slots: torch.Tensor, + kv_caches: list[tuple[torch.Tensor | None, ...]], + layer_synchronizer: LayerSynchronizer | None, + ) -> torch.Tensor | None: + return self.model.write_context_kv( + target_hidden, + positions, + cache_slots, + kv_caches, + layer_synchronizer, + ) diff --git a/xllm/python/registry.py b/xllm/python/registry.py index fd257a40f5..79da198365 100644 --- a/xllm/python/registry.py +++ b/xllm/python/registry.py @@ -91,6 +91,18 @@ def _register_builtin_models() -> None: "qwen3_5_moe", "qwen3_5_moe_text", ) + _register_model_path( + "xllm.python.models.qwen3_dspark", + "Qwen3DSparkForCausalLM", + "DSparkDraftModel", + "Qwen3DSparkModel", + ) + _register_model_path( + "xllm.python.models.qwen3_dflash", + "DFlashQwen3ForCausalLM", + "DFlashDraftModel", + "DFlashQwen3Model", + ) _register_model_path( "xllm.python.models.qwen3_vl", "Qwen3VLForConditionalGeneration",