feat: support qwen3 dspark decoding on NPU with Python executor.

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);

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()

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)
 
 

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]]))