diff --git a/tests/core/kernels/npu/npu_xllm_ops_test.cpp b/tests/core/kernels/npu/npu_xllm_ops_test.cpp index 3a47ab700d..5d026517e3 100644 --- a/tests/core/kernels/npu/npu_xllm_ops_test.cpp +++ b/tests/core/kernels/npu/npu_xllm_ops_test.cpp @@ -209,6 +209,425 @@ TEST_F(NpuXllmOpsTest, EmbeddedInterpreterSeesOps) { .item(); } +TEST_F(NpuXllmOpsTest, GroupGemmWrapperMatchesInt32Reference) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch + +device = torch.device("privateuseone:0") +torch.manual_seed(20260814) +tokens, experts, input_dim, output_dim = 8, 2, 128, 256 +x_cpu = torch.randint(-4, 5, (tokens, input_dim), dtype=torch.int8) +w_cpu = torch.randint(-4, 5, (experts, input_dim, output_dim), dtype=torch.int8) +group_list_cpu = torch.tensor([4, 4], dtype=torch.int64) + +x = x_cpu.to(device) +w = w_cpu.to(device) +group_list = group_list_cpu.to(device) +out = torch.ops.xllm_ops.group_gemm( + x, w, None, None, group_list, 2, 0, 1, torch.int32 +) +torch.npu.synchronize() + +expected = torch.cat(( + x_cpu[:4].to(torch.int32) @ w_cpu[0].to(torch.int32), + x_cpu[4:].to(torch.int32) @ w_cpu[1].to(torch.int32), +), dim=0) +assert out.shape == (tokens, output_dim) +assert out.dtype == torch.int32 +torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) +)PY"); +} + +TEST_F(NpuXllmOpsTest, GroupGemmWrapperAcceptsScaleAndPerTokenScale) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch + +device = torch.device("privateuseone:0") +tokens, experts, input_dim, output_dim = 8, 2, 128, 128 +x = torch.randint(-4, 5, (tokens, input_dim), dtype=torch.int8, device=device) +w = torch.randint(-4, 5, (experts, input_dim, output_dim), dtype=torch.int8, device=device) +scale = torch.ones((experts, output_dim), dtype=torch.bfloat16, device=device) +per_token_scale = torch.ones((tokens,), dtype=torch.float32, device=device) +group_list = torch.tensor([4, 4], dtype=torch.int64, device=device) + +out = torch.ops.xllm_ops.group_gemm( + x, w, scale, per_token_scale, group_list, 2, 0, 1, torch.bfloat16 +) +torch.npu.synchronize() +assert out.shape == (tokens, output_dim) +assert out.dtype == torch.bfloat16 +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4PartialRotaryPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.rotary_embedding import ( + npu_inplace_partial_rotary_mul, +) + +torch.manual_seed(2026) +x_cpu = torch.randn((8, 2, 128), dtype=torch.float32).to(torch.bfloat16) +cos_cpu = torch.randn((8, 64), dtype=torch.float32).to(torch.bfloat16) +sin_cpu = torch.randn((8, 64), dtype=torch.float32).to(torch.bfloat16) + +expected = x_cpu.float().clone() +segment = x_cpu[..., 64:128].float() +swapped = torch.empty_like(segment) +swapped[..., 0::2] = segment[..., 1::2] +swapped[..., 1::2] = segment[..., 0::2] +sign = torch.ones_like(cos_cpu.float()) +sign[..., 0::2] = -1 +expected[..., 64:128] = ( + segment * cos_cpu.float().unsqueeze(1) + + swapped * sin_cpu.float().unsqueeze(1) * sign.unsqueeze(1) +) +expected = expected.to(torch.bfloat16).float() + +x = x_cpu.to("privateuseone:0") +cos = cos_cpu.to(x.device) +sin = sin_cpu.to(x.device) +result = npu_inplace_partial_rotary_mul(x, cos, sin, 64, 64) +torch.npu.synchronize() + +assert result.data_ptr() == x.data_ptr() +torch.testing.assert_close( + x.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4CompressorPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import compressor + +device = torch.device("privateuseone:0") +torch.manual_seed(2025) +batch, tokens, hidden = 1, 128, 1024 +ratio, head_dim, coff, rope_dim = 128, 512, 1, 64 +compressed_tokens = tokens // ratio + +x_cpu = (torch.randn(batch, tokens, hidden) * 0.1).to(torch.float16) +wkv_cpu = (torch.randn(coff * head_dim, hidden) * 0.05).to(torch.float16) +wgate_cpu = (torch.randn(coff * head_dim, hidden) * 0.05).to(torch.float16) +ape_cpu = (torch.randn(ratio, coff * head_dim) * 0.1).float() +norm_cpu = (torch.randn(head_dim) * 0.1 + 1).to(torch.float16) +rope_cos_cpu = ( + torch.randn(batch, compressed_tokens, rope_dim) * 0.1 +).to(torch.float16) +rope_sin_cpu = ( + torch.randn(batch, compressed_tokens, rope_dim) * 0.1 +).to(torch.float16) + +projected_kv = x_cpu.float()[0] @ wkv_cpu.float().T +scores = x_cpu.float()[0] @ wgate_cpu.float().T + ape_cpu +pooled = (torch.softmax(scores, dim=0) * projected_kv).sum(0, keepdim=True) +variance = pooled.square().mean(-1, keepdim=True) +expected = pooled * torch.rsqrt(variance + 1e-6) * norm_cpu.float() +rope_segment = expected[:, -rope_dim:].clone() +half = rope_dim // 2 +rotated = torch.cat((-rope_segment[:, half:], rope_segment[:, :half]), dim=-1) +expected[:, -rope_dim:] = ( + rope_segment * rope_cos_cpu.float()[0] + + rotated * rope_sin_cpu.float()[0] +) +expected = expected.view(batch, compressed_tokens, head_dim).half().float() + +x = x_cpu.to(device) +wkv = wkv_cpu.to(device) +wgate = wgate_cpu.to(device) +ape = ape_cpu.to(device) +norm_weight = norm_cpu.to(device) +rope_sin = rope_sin_cpu.to(device) +rope_cos = rope_cos_cpu.to(device) +kv_state = torch.zeros((1, 128, head_dim), dtype=torch.float32, device=device) +score_state = torch.zeros_like(kv_state) +kv_block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +score_block_table = torch.tensor([[0]], dtype=torch.int32, device=device) + +out, wkv_proj, softmax_res, norm_x, norm_rstd = compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table, + score_block_table, + None, + None, + None, + rope_dim, + ratio, + coff, + 1e-6, + 1, + False, +) +torch.npu.synchronize() + +assert out.shape == (batch, compressed_tokens, head_dim) +assert out.dtype == torch.float16 +assert wkv_proj.numel() == 0 +assert softmax_res.numel() == 0 +assert norm_x.numel() == 0 +assert norm_rstd.numel() == 0 +torch.testing.assert_close( + out.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4QuantLightningIndexerPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import ( + quant_lightning_indexer, + quant_lightning_indexer_metadata, +) + +device = torch.device("privateuseone:0") +torch.manual_seed(2026) +tokens, heads, head_dim = 84, 64, 128 +page_size, sparse_count = 128, 512 +query_cpu = torch.randint(-8, 8, (tokens, heads, head_dim), dtype=torch.int8) +key_cpu = torch.randint(-8, 8, (1, page_size, 1, head_dim), dtype=torch.int8) +query = query_cpu.to(device) +key = key_cpu.to(device) +weights = torch.ones((tokens, heads), dtype=torch.float16, device=device) +query_scale = torch.ones((tokens, heads), dtype=torch.float16, device=device) +key_scale = torch.ones((1, page_size, 1), dtype=torch.float16, device=device) +query_lens = torch.tensor([tokens], dtype=torch.int32, device=device) +key_lens = torch.tensor([tokens], dtype=torch.int32, device=device) +block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +metadata = quant_lightning_indexer_metadata( + heads, + 1, + head_dim, + 0, + 0, + query_lens, + key_lens, + 1, + tokens, + tokens, + "TND", + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + 4, + "npu", +) +metadata_again = quant_lightning_indexer_metadata( + heads, + 1, + head_dim, + 0, + 0, + query_lens, + key_lens, + 1, + tokens, + tokens, + "TND", + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + 4, + "npu", +) +indices, values = quant_lightning_indexer( + query, + key, + weights, + query_scale, + key_scale, + 0, + 0, + query_lens, + key_lens, + block_table, + metadata, + "TND", + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + 4, + False, +) +torch.npu.synchronize() + +assert indices.shape == (tokens, 1, sparse_count) +assert indices.dtype == torch.int32 +assert values.numel() == 0 +assert values.dtype == torch.float32 +assert torch.equal(metadata.cpu(), metadata_again.cpu()) + +valid_key_count = tokens // 4 +indices_cpu = indices.cpu().squeeze(1) +assert torch.all( + (indices_cpu == -1) + | ((indices_cpu >= 0) & (indices_cpu < valid_key_count)) +) +keys = key_cpu[0, :valid_key_count, 0].float() +token_idx = tokens - 1 +dots = query_cpu[token_idx].float() @ keys.T +expected_top8 = set(torch.topk(dots.clamp_min(0).sum(0), 8).indices.tolist()) +actual_top8 = set(indices_cpu[token_idx, :8].tolist()) +assert len(expected_top8 & actual_top8) >= 4, ( + sorted(expected_top8), + sorted(actual_top8), +) +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4SparseAttentionPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import ( + sparse_attn_sharedkv, + sparse_attn_sharedkv_metadata, +) + +device = torch.device("privateuseone:0") +torch.manual_seed(1234) +batch, q_tokens, kv_tokens = 1, 4, 16 +heads, head_dim, page_size = 64, 512, 16 +query_cpu = (torch.randn(batch, q_tokens, heads, head_dim) * 0.1).half() +kv_cpu = (torch.randn(batch, kv_tokens, 1, head_dim) * 0.1).half() +sinks_cpu = (torch.randn(heads) * 0.1).float() +query = query_cpu.to(device) +ori_kv = kv_cpu.view(1, page_size, 1, head_dim).to(device) +block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +cu_q = torch.tensor([0, q_tokens], dtype=torch.int32, device=device) +cu_kv = torch.tensor([0, kv_tokens], dtype=torch.int32, device=device) +seq_q = torch.tensor([q_tokens], dtype=torch.int32, device=device) +seq_kv = torch.tensor([kv_tokens], dtype=torch.int32, device=device) +sinks = sinks_cpu.to(device) +metadata = sparse_attn_sharedkv_metadata( + heads, + 1, + head_dim, + cu_q, + cu_kv, + None, + seq_q, + seq_kv, + batch, + q_tokens, + kv_tokens, + 0, + 0, + 1, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + True, + False, +) +metadata_again = sparse_attn_sharedkv_metadata( + heads, + 1, + head_dim, + cu_q, + cu_kv, + None, + seq_q, + seq_kv, + batch, + q_tokens, + kv_tokens, + 0, + 0, + 1, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + True, + False, +) +out, lse = sparse_attn_sharedkv( + query, + ori_kv, + None, + None, + None, + block_table, + None, + None, + None, + None, + None, + seq_kv, + sinks, + metadata, + head_dim**-0.5, + 1, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + False, +) +torch.npu.synchronize() + +assert out.shape == query.shape +assert out.dtype == query.dtype +assert lse.numel() == 0 +assert torch.equal(metadata.cpu(), metadata_again.cpu()) + +expected = torch.zeros_like(query_cpu.float()) +keys = kv_cpu[0, :, 0].float() +scale = head_dim**-0.5 +for q_idx in range(q_tokens): + diagonal = kv_tokens - q_tokens + q_idx + left = max(diagonal - 127, 0) + right = diagonal + selected_keys = keys[left:right + 1] + logits = query_cpu[0, q_idx].float() @ selected_keys.T * scale + sink_logits = sinks_cpu[:, None] + normalizer = torch.logsumexp( + torch.cat((logits, sink_logits), dim=1), dim=1 + ) + probabilities = torch.exp(logits - normalizer[:, None]) + expected[0, q_idx] = probabilities @ selected_keys +expected = expected.half().float() +torch.testing.assert_close( + out.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + TEST_F(NpuXllmOpsTest, Qwen35_27B_TP4_FullAttentionMatchesReference) { py::gil_scoped_acquire gil; if (!is_ascend950_device()) { diff --git a/tests/python/test_collectives.py b/tests/python/test_collectives.py index fbef228c57..65fbe26ba7 100644 --- a/tests/python/test_collectives.py +++ b/tests/python/test_collectives.py @@ -18,20 +18,17 @@ import importlib.util import json +import sys from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock import pytest import torch import torch.distributed as dist - -_MODULE_PATH = ( - Path(__file__).parents[2] / "xllm" / "python" / "distributed" / "collectives.py" -) -_SPEC = importlib.util.spec_from_file_location( - "_xllm_collectives_under_test", _MODULE_PATH -) +_MODULE_PATH = Path(__file__).parents[2] / "xllm" / "python" / "distributed" / "collectives.py" +_SPEC = importlib.util.spec_from_file_location("_xllm_collectives_under_test", _MODULE_PATH) assert _SPEC is not None and _SPEC.loader is not None collectives = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(collectives) @@ -69,9 +66,7 @@ def __init__(self, topology: list[dict[str, object]] | None = None) -> None: self.values: dict[str, bytes] = {} if topology is not None: for rank, entry in enumerate(topology): - self.values[ - f"xllm/python_collectives/topology/v1/{rank}" - ] = json.dumps(entry).encode("utf-8") + self.values[f"xllm/python_collectives/topology/v1/{rank}"] = json.dumps(entry).encode("utf-8") def set(self, key: str, value: str) -> None: self.values[key] = value.encode("utf-8") @@ -91,9 +86,7 @@ def _mock_process_groups( handed, which is what the module checks its caller's rank against. """ if topology is None: - topology = [ - {"hostname": "node-0", "device_index": rank} for rank in range(16) - ] + topology = [{"hostname": "node-0", "device_index": rank} for rank in range(16)] base_store = _FakeStore(topology) tcp_store = MagicMock(return_value=base_store) init_world = MagicMock() @@ -111,16 +104,10 @@ def _mock_process_groups( def test_parallel_groups_share_one_multitenant_tcp_store(monkeypatch): - base_store, tcp_store, init_world, new_group = _mock_process_groups( - monkeypatch, global_rank=0 - ) + base_store, tcp_store, init_world, new_group = _mock_process_groups(monkeypatch, global_rank=0) - collectives.init_process_group( - "tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 0, 2, 0 - ) - collectives.init_process_group( - "moe_tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 0, 2, 0 - ) + collectives.init_process_group("tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 0, 2, 0) + collectives.init_process_group("moe_tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 0, 2, 0) tcp_store.assert_called_once() assert tcp_store.call_args.args[:4] == ("127.0.0.1", 46001, 2, True) @@ -139,12 +126,47 @@ def test_parallel_groups_share_one_multitenant_tcp_store(monkeypatch): ] +def test_native_runtime_bridge_bypasses_python_process_groups(monkeypatch): + calls: list[str] = [] + + runtime = SimpleNamespace( + tp_all_reduce=lambda tensor: (calls.append("tp_reduce"), tensor.add_(1)), + tp_all_gather=lambda tensor, dim: ( + calls.append(f"tp_gather:{dim}"), + torch.cat((tensor, tensor), dim=dim), + )[1], + moe_tp_all_reduce=lambda tensor: ( + calls.append("moe_tp_reduce"), + tensor.add_(2), + ), + moe_ep_all_reduce=lambda tensor: ( + calls.append("moe_ep_reduce"), + tensor.add_(4), + ), + ) + monkeypatch.setitem(sys.modules, "xllm_runtime", runtime) + python_reduce = MagicMock(side_effect=AssertionError("c10d fallback used")) + python_gather = MagicMock(side_effect=AssertionError("c10d fallback used")) + monkeypatch.setattr(collectives, "all_reduce_", python_reduce) + monkeypatch.setattr(collectives, "all_gather", python_gather) + + value = torch.tensor([[1.0]]) + collectives.tp_all_reduce(value) + gathered = collectives.tp_all_gather(value, 1, 2) + collectives.moe_tp_all_reduce(value) + collectives.moe_ep_all_reduce(value) + + assert calls == ["tp_reduce", "tp_gather:1", "moe_tp_reduce", "moe_ep_reduce"] + assert gathered.tolist() == [[2.0, 2.0]] + assert value.tolist() == [[8.0]] + python_reduce.assert_not_called() + python_gather.assert_not_called() + + def test_tcp_store_master_is_global_rank_zero_not_group_rank_zero(monkeypatch): _, tcp_store, _, _ = _mock_process_groups(monkeypatch, global_rank=2) - collectives.init_process_group( - "tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 2, 4, 1 - ) + collectives.init_process_group("tp", "127.0.0.1", 46001, 0, 2, "cuda:0", 2, 4, 1) assert tcp_store.call_args.args[:4] == ("127.0.0.1", 46001, 4, False) @@ -157,9 +179,7 @@ def test_symmetric_memory_rejects_cross_host_group(monkeypatch): can_access_peer = MagicMock(return_value=True) monkeypatch.setattr(torch.cuda, "can_device_access_peer", can_access_peer) - assert not collectives._supports_symmetric_memory( - torch.device("cuda:0"), [0, 1] - ) + assert not collectives._supports_symmetric_memory(torch.device("cuda:0"), [0, 1]) can_access_peer.assert_not_called() @@ -174,9 +194,7 @@ def test_symmetric_memory_rejects_incomplete_peer_domain(monkeypatch): lambda source, destination: (source, destination) != (1, 0), ) - assert not collectives._supports_symmetric_memory( - torch.device("cuda:0"), [0, 1] - ) + assert not collectives._supports_symmetric_memory(torch.device("cuda:0"), [0, 1]) @pytest.mark.parametrize("dtype", [torch.float16, torch.float64, torch.int32]) diff --git a/tests/python/test_deepseek_v4_model.py b/tests/python/test_deepseek_v4_model.py new file mode 100644 index 0000000000..55e50b2588 --- /dev/null +++ b/tests/python/test_deepseek_v4_model.py @@ -0,0 +1,342 @@ +# 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. + +"""DeepSeek-V4 Python model: config parsing, registry, structure. + +Pure-Python: does not load compiled operators or weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from xllm.python.models import deepseek_v4, deepseek_v32 +from xllm.python.models.deepseek_v4 import ( + DeepseekV4Config, + DeepseekV4DecoderLayer, + DeepseekV4HyperConnection, + DeepseekV4Model, + DeepseekV4MoE, + DeepseekV4RotaryEmbedding, +) +from xllm.python.models.deepseek_v32 import W8A8DynamicLinear, _swiglu_with_clamp +from xllm.python.registry import get_model_class + +_DSV4_CONFIG = { + "model_type": "deepseek_v4", + "architectures": ["DeepseekV4ForCausalLM"], + "hidden_size": 4096, + "num_hidden_layers": 4, + "num_attention_heads": 64, + "head_dim": 512, + "vocab_size": 129280, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "max_position_embeddings": 1048576, + "original_max_position_embeddings": 65536, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 16, + "original_max_position_embeddings": 65536, + "type": "yarn", + }, + "q_lora_rank": 1024, + "qk_rope_head_dim": 64, + "o_lora_rank": 1024, + "o_groups": 8, + "compress_ratios": [0, 4, 128, 4], + "window_size": 128, + "sliding_window": 128, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 512, + "n_activated_experts": 6, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-6, + "scoring_func": "sqrtsoftplus", + "scale_fmt": "ue8m0", + "n_routed_experts": 256, + "moe_intermediate_size": 2048, + "first_k_dense_replace": 0, + "tie_word_embeddings": False, +} + + +def test_config_from_dict_reads_dsv4_fields() -> None: + cfg = DeepseekV4Config.from_dict(_DSV4_CONFIG) + assert cfg.model_type == "deepseek_v4" + assert cfg.n_layers == 4 + assert cfg.compress_ratios == [1, 4, 128, 4] + assert cfg.window_size == 128 + assert cfg.o_lora_rank == 1024 + assert cfg.o_groups == 8 + assert cfg.hc_mult == 4 + assert cfg.index_topk == 512 + assert cfg.rope_scaling_factor == 16.0 + + +def test_config_prefers_dsv4_model_args_over_zero_legacy_rope_fields() -> None: + cfg = DeepseekV4Config.from_dict( + { + **_DSV4_CONFIG, + "rope_scaling": None, + "factor": 16.0, + "beta_fast": 32.0, + "beta_slow": 1.0, + "rope_scaling_attn_factor": 1.0, + "rope_scaling_factor": 0.0, + "rope_scaling_beta_fast": 0.0, + "rope_scaling_beta_slow": 0.0, + } + ) + + assert cfg.rope_scaling_factor == 16.0 + assert cfg.rope_beta_fast == 32 + assert cfg.rope_beta_slow == 1 + assert cfg.rope_mscale == 1.0 + + +def test_rotary_cache_matches_cpp_cpu_float32_construction() -> None: + rotary = DeepseekV4RotaryEmbedding( + rotary_dim=64, + max_position_embeddings=87, + scaling_factor=16.0, + theta=10000.0, + beta_fast=32, + beta_slow=1, + old_context_len=1048576, + dtype=torch.bfloat16, + device=torch.device("cpu"), + ) + + # Position 86, frequency 1 is the first observed CPU/NPU rounding split. + # Lock the Python cache to the value produced by the C++ CPU path. + assert rotary.cos_sin_cache[86, 1].item() == -0.087890625 + + +def test_rotary_cache_shares_identical_descriptors() -> None: + args = dict( + rotary_dim=64, + max_position_embeddings=87, + scaling_factor=16.0, + theta=160000.0, + beta_fast=32, + beta_slow=1, + old_context_len=1048576, + dtype=torch.bfloat16, + device=torch.device("cpu"), + ) + c4 = DeepseekV4RotaryEmbedding(**args) + c128 = DeepseekV4RotaryEmbedding(**args) + default = DeepseekV4RotaryEmbedding(**{**args, "theta": 10000.0}) + + assert c4.cos_sin_cache.data_ptr() == c128.cos_sin_cache.data_ptr() + assert c4.cos_sin_cache.data_ptr() != default.cos_sin_cache.data_ptr() + + +def test_registry_resolves_deepseek_v4() -> None: + cls = get_model_class("deepseek_v4") + assert cls.__name__ == "DeepseekV4ForCausalLM" + + +def test_hyper_connection_shapes() -> None: + cfg = DeepseekV4Config.from_dict(_DSV4_CONFIG) + hc = DeepseekV4HyperConnection(cfg, torch.float32, torch.device("cpu")) + assert hc.hc_mult_local == 4 # hc_mult is NOT TP-sharded + # hc_*_fn: [mix_hc, hc_dim] = [(2+mult)*mult, mult*hidden] = [24, 16384]. + assert hc.hc_attn_fn.shape == ((2 + 4) * 4, 4 * 4096) + assert hc.hc_attn_base.shape == ((2 + 4) * 4,) + assert hc.hc_attn_scale.shape == (3,) + # hc_pre calls the compiled NPU kernel, which is not available in the + # pure-Python unit-test context; only shape construction is verified here. + + +def test_decoder_layer_builds() -> None: + """A C4 decoder layer builds attention + HC + MoE without error.""" + cfg = DeepseekV4Config.from_dict(_DSV4_CONFIG) + layer = DeepseekV4DecoderLayer(cfg, layer_id=1, dtype=torch.float32, device=torch.device("cpu")) + assert layer.self_attn.layer_id == 1 + assert layer.self_attn.indexer is not None + assert layer.hc.hc_mult_local == 4 + + +def test_attention_builds_compression_modules_only_for_matching_ratios() -> None: + cfg = DeepseekV4Config.from_dict(_DSV4_CONFIG) + c1 = DeepseekV4DecoderLayer(cfg, layer_id=0, dtype=torch.float32, device=torch.device("cpu")).self_attn + c4 = DeepseekV4DecoderLayer(cfg, layer_id=1, dtype=torch.float32, device=torch.device("cpu")).self_attn + c128 = DeepseekV4DecoderLayer(cfg, layer_id=2, dtype=torch.float32, device=torch.device("cpu")).self_attn + + assert c1.indexer is None + assert not hasattr(c1, "cmp_wkv") + assert c4.indexer is not None + assert hasattr(c4, "cmp_wkv") + assert c128.indexer is None + assert hasattr(c128, "cmp_wkv") + + +def test_moe_gate_state_matches_cpp_parameter_ownership() -> None: + cfg_dict = dict(_DSV4_CONFIG) + cfg_dict["n_hash_layers"] = 3 + cfg = DeepseekV4Config.from_dict(cfg_dict) + + hash_moe = DeepseekV4MoE(cfg, layer_id=2, dtype=torch.float32, device=torch.device("cpu")) + non_hash_moe = DeepseekV4MoE(cfg, layer_id=3, dtype=torch.float32, device=torch.device("cpu")) + + hash_params = dict(hash_moe.named_parameters()) + non_hash_params = dict(non_hash_moe.named_parameters()) + assert "tid2eid" in hash_params + assert not hash_params["tid2eid"].requires_grad + assert "e_score_correction_bias" in non_hash_params + assert not non_hash_params["e_score_correction_bias"].requires_grad + + +def test_clamped_swiglu_matches_cpp_activation_formula() -> None: + x = torch.tensor([[-20.0, 5.0, 20.0, 12.0, -15.0, 3.0]], dtype=torch.bfloat16) + gate, up = x.chunk(2, dim=-1) + expected = (torch.nn.functional.silu(gate.float().clamp_max(10.0)) * up.float().clamp(min=-10.0, max=10.0)).to( + x.dtype + ) + + torch.testing.assert_close(_swiglu_with_clamp(x, 10.0), expected) + + +def test_dynamic_linear_preserves_v3_and_v4_weight_layout_contracts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[bool] = [] + + def fake_quant_matmul(x, weight, transpose2, scale, offset, pertoken, bias, output_dtype): + calls.append(transpose2) + return torch.empty((x.size(0), scale.numel()), dtype=output_dtype) + + monkeypatch.setattr(deepseek_v32.kernels, "quant_matmul", fake_quant_matmul, raising=False) + x = torch.ones((2, 3), dtype=torch.int8) + pertoken = torch.ones((2,), dtype=torch.float32) + + v3 = W8A8DynamicLinear(3, 4, torch.device("cpu")) + v3.weight_offset.zero_() + v3.process_weights_after_loading() + assert v3.weight.shape == (3, 4) + v3.forward_quant(x, pertoken) + + v4 = W8A8DynamicLinear(3, 4, torch.device("cpu"), transpose_weight_after_loading=False) + v4.weight_offset.zero_() + v4.process_weights_after_loading() + assert v4.weight.shape == (4, 3) + v4.forward_quant(x, pertoken) + + assert calls == [False, True] + + +def test_model_rejects_cp_until_cp_context_is_available() -> None: + cfg = DeepseekV4Config.from_dict({**_DSV4_CONFIG, "cp_size": 2}) + with pytest.raises(NotImplementedError, match="CP context PR"): + DeepseekV4Model(cfg, torch.float32, torch.device("cpu")) + + +def test_moe_uses_dedicated_group_sizes() -> None: + cfg_dict = dict(_DSV4_CONFIG) + cfg_dict.update( + tp_size=4, + tp_rank=1, + ep_size=8, + ep_rank=3, + moe_tp_size=1, + moe_tp_rank=0, + cp_size=1, + cp_rank=0, + ) + cfg = DeepseekV4Config.from_dict(cfg_dict) + moe = DeepseekV4MoE(cfg, layer_id=0, dtype=torch.float32, device=torch.device("cpu")) + + assert moe.moe_tp_size == 1 + assert moe.moe_tp_rank == 0 + assert moe.start_expert_id == 3 * (cfg.n_routed_experts // cfg.ep_size) + + +def test_moe_tp_ep_reduction_order_matches_cpp(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + class FakeDistributed: + @staticmethod + def moe_tp_all_reduce(tensor: torch.Tensor) -> None: + calls.append("moe_tp") + tensor.add_(10) + + @staticmethod + def moe_ep_all_reduce(tensor: torch.Tensor) -> None: + calls.append("moe_ep") + tensor.add_(100) + + monkeypatch.setattr(deepseek_v4, "distributed", FakeDistributed) + owner = SimpleNamespace( + cfg=SimpleNamespace(ep_size=2, tp_size=1), + moe_tp_size=2, + ) + routed = torch.zeros(2) + shared = torch.ones(2) + + output = DeepseekV4MoE._reduce_moe_outputs(owner, routed, shared) + + assert calls == ["moe_tp", "moe_ep", "moe_tp"] + assert torch.equal(output, torch.full((2,), 121.0)) + + +def test_moe_ep_only_reduces_routed_output(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + class FakeDistributed: + @staticmethod + def moe_ep_all_reduce(tensor: torch.Tensor) -> None: + calls.append("moe_ep") + tensor.add_(100) + + monkeypatch.setattr(deepseek_v4, "distributed", FakeDistributed) + owner = SimpleNamespace( + cfg=SimpleNamespace(ep_size=2, tp_size=1), + moe_tp_size=1, + ) + + output = DeepseekV4MoE._reduce_moe_outputs(owner, torch.zeros(1), torch.ones(1)) + + assert calls == ["moe_ep"] + assert torch.equal(output, torch.full((1,), 101.0)) + + +def test_moe_tp_only_combines_before_one_reduce( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + class FakeDistributed: + @staticmethod + def moe_tp_all_reduce(tensor: torch.Tensor) -> None: + calls.append("moe_tp") + tensor.mul_(2) + + monkeypatch.setattr(deepseek_v4, "distributed", FakeDistributed) + owner = SimpleNamespace( + cfg=SimpleNamespace(ep_size=1, tp_size=2), + moe_tp_size=2, + ) + + output = DeepseekV4MoE._reduce_moe_outputs(owner, torch.full((1,), 2.0), torch.full((1,), 3.0)) + + assert calls == ["moe_tp"] + assert torch.equal(output, torch.full((1,), 10.0)) diff --git a/tests/python/test_dsa_attention.py b/tests/python/test_dsa_attention.py new file mode 100644 index 0000000000..b299fada22 --- /dev/null +++ b/tests/python/test_dsa_attention.py @@ -0,0 +1,378 @@ +# 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. + +"""Unit tests for the DSA attention backend (cache-mapping + slot scatter). + +Pure-Python: does not load compiled operators. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from xllm.python import kernels +from xllm.python.attention import dsa_attention as dsa_attention_module +from xllm.python.attention.backend import LayerCache +from xllm.python.attention.dsa_attention import ( + DsaAttentionBackend, + _DsaCacheMapping, + _get_layer_cache_tensor, + _scatter_by_slot, +) +from xllm.python.attention.dsa_metadata import build_cache_specs + + +def _make_backend() -> DsaAttentionBackend: + compress_ratios = [0, 4, 128] + caches_info, group_infos = build_cache_specs(compress_ratios, 128, 3) + return DsaAttentionBackend( + compress_ratios=compress_ratios, + window_size=128, + n_layers=3, + num_heads=8, + attn_head_dim=512, + index_topk=512, + index_n_heads=64, + index_head_dim=128, + rope_head_dim=64, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + +def test_compress_ratio_per_layer() -> None: + b = _make_backend() + assert b._layer_compress_ratio(0) == 1 + assert b._layer_compress_ratio(1) == 4 + assert b._layer_compress_ratio(2) == 128 + + +def test_resolve_cache_mapping_c4() -> None: + """C4 layer: cmp/index/indexer_scale (TOKEN, idx 0/1/7), ori..index_score (SWA, 2..6).""" + b = _make_backend() + m = b._resolve_cache_mapping(1, 4) + assert m.cmp_cache_idx == 0 + assert m.index_cache_idx == 1 + assert m.indexer_scale_cache_idx == 7 + assert m.ori_cache_idx == 2 + assert m.kv_state_cache_idx == 3 + assert m.score_state_cache_idx == 4 + assert m.index_kv_state_cache_idx == 5 + assert m.index_score_state_cache_idx == 6 + + +def test_resolve_cache_mapping_c128() -> None: + """C128 layer: cmp (TOKEN, idx 0), ori..index_score (SWA, 1..3).""" + b = _make_backend() + m = b._resolve_cache_mapping(2, 128) + assert m.cmp_cache_idx == 0 + assert m.ori_cache_idx == 1 + assert m.kv_state_cache_idx == 2 + assert m.score_state_cache_idx == 3 + assert m.index_cache_idx == -1 # C128 has no indexer + assert m.indexer_scale_cache_idx == -1 + + +def test_resolve_cache_mapping_c1() -> None: + """C1 layer: only one SWA cache; compress_ratio==1 -> no cmp.""" + b = _make_backend() + m = b._resolve_cache_mapping(0, 1) + assert m.cmp_cache_idx == -1 + assert m.ori_cache_idx == 0 + assert m.index_cache_idx == -1 + + +def test_get_layer_cache_tensor_bounds() -> None: + tensors = [[torch.empty(0)], [torch.zeros(2), torch.zeros(3)]] + assert _get_layer_cache_tensor(tensors, 0, 0).numel() == 0 + assert _get_layer_cache_tensor(tensors, 1, 1).numel() == 3 + assert _get_layer_cache_tensor(tensors, 5, 0) is None # bad layer + assert _get_layer_cache_tensor(tensors, 1, 9) is None # bad cache idx + + +def test_scatter_by_slot_writes_rows() -> None: + cache = torch.zeros(4, 3, dtype=torch.float32) + # slot 0 -> row 0, slot 2 -> row 2, slot -1 skipped. + slots = torch.tensor([0, -1, 2], dtype=torch.int32) + value = torch.tensor([[1.0, 1.0, 1.0], [9.0, 9.0, 9.0], [2.0, 2.0, 2.0]]) + _scatter_by_slot(cache, slots, value) + assert torch.equal(cache[0], torch.tensor([1.0, 1.0, 1.0])) + assert torch.equal(cache[2], torch.tensor([2.0, 2.0, 2.0])) + # Row 1 untouched (slot -1 skipped); value row 1 dropped. + assert torch.equal(cache[1], torch.zeros(3)) + + +def test_scatter_by_slot_ignores_all_padded_rows() -> None: + cache = torch.arange(12, dtype=torch.float32).view(4, 3) + original = cache.clone() + slots = torch.full((2,), -1, dtype=torch.int32) + values = torch.full((2, 3), 99.0) + + _scatter_by_slot(cache, slots, values) + + assert torch.equal(cache, original) + + +def test_default_mapping_is_empty() -> None: + m = _DsaCacheMapping() + assert m.cmp_cache_idx == -1 + assert m.ori_cache_idx == -1 + + +def test_prepare_binds_dsa_metadata_to_current_forward(monkeypatch) -> None: + backend = _make_backend() + monkeypatch.setattr(backend, "_move_metadata_to_device", lambda dsa: None) + monkeypatch.setattr(backend, "_build_precomputed_metadata", lambda dsa, metadata: None) + + def make_metadata(kv_len: int, is_prefill: bool) -> SimpleNamespace: + q_len = kv_len if is_prefill else 1 + return SimpleNamespace( + multi_block_tables=[], + kv_seq_lens_host=torch.tensor([kv_len], dtype=torch.int32), + q_seq_lens_host=torch.tensor([q_len], dtype=torch.int32), + is_prefill=is_prefill, + is_chunked_prefill=False, + dsa_metadata=None, + dsa_positions=None, + dsa_cos_sin=None, + dsa_c4_cos_sin=None, + dsa_c128_cos_sin=None, + dsa_graph_block_table_cols=0, + dsa_graph_mode=False, + ) + + prefill = make_metadata(84, True) + backend.prepare(prefill) + backend.prepare_dsa_metadata_for_forward() + prefill_dsa = prefill.dsa_metadata + assert prefill_dsa.max_query_len == 84 + + decode = make_metadata(85, False) + backend.prepare(decode) + backend.prepare_dsa_metadata_for_forward() + assert decode.dsa_metadata is not prefill_dsa + assert decode.dsa_metadata.max_query_len == 1 + assert prefill.dsa_metadata is prefill_dsa + + +def test_graph_mode_is_explicitly_deferred() -> None: + backend = _make_backend() + with pytest.raises(NotImplementedError, match="ACL graph"): + backend.prepare(SimpleNamespace(), graph_mode=True) + + +def test_decode_precomputed_metadata_matches_cpp_contract(monkeypatch) -> None: + backend = _make_backend() + sparse_calls: list[dict] = [] + qli_calls: list[dict] = [] + + def fake_sparse_metadata(**kwargs): + sparse_calls.append(kwargs) + return torch.tensor([kwargs["cmp_ratio"]], dtype=torch.int32) + + def fake_qli_metadata(**kwargs): + qli_calls.append(kwargs) + return torch.tensor([4], dtype=torch.int32) + + monkeypatch.setattr( + kernels, + "sparse_attn_sharedkv_metadata", + fake_sparse_metadata, + raising=False, + ) + monkeypatch.setattr( + kernels, + "quant_lightning_indexer_metadata", + fake_qli_metadata, + raising=False, + ) + dsa = SimpleNamespace( + actual_seq_lengths_query=torch.tensor([0, 1], dtype=torch.int32), + actual_seq_lengths_kv=torch.tensor([85], dtype=torch.int32), + seq_lens_q=torch.tensor([1], dtype=torch.int32), + seq_lens=torch.tensor([85], dtype=torch.int32), + max_query_len=1, + max_seq_len=85, + ) + metadata = SimpleNamespace( + max_query_len=1, + max_seq_len=85, + q_seq_lens_host=torch.tensor([1], dtype=torch.int32), + kv_seq_lens_host=torch.tensor([85], dtype=torch.int32), + ) + + backend._build_precomputed_metadata(dsa, metadata) + + assert [call["cmp_ratio"] for call in sparse_calls] == [1, 4, 128] + assert all(call["head_dim"] == 512 for call in sparse_calls) + assert all(call["cu_seqlens_q"].tolist() == [0, 1] for call in sparse_calls) + assert all(call["cu_seqlens_ori_kv"].numel() == 0 for call in sparse_calls) + assert sparse_calls[1]["cmp_topk"] == 512 + assert qli_calls[0]["actual_seq_lengths_query"].tolist() == [1] + assert qli_calls[0]["actual_seq_lengths_key"].tolist() == [85] + assert qli_calls[0]["head_dim"] == 128 + assert dsa.precomputed_metadata_inputs[0] is dsa.actual_seq_lengths_query + + +def test_c4_execute_requires_model_compressor(monkeypatch) -> None: + backend = _make_backend() + empty_cache = LayerCache(key=None, value=None) + cmp_cache = torch.zeros(2, 128, 1, 512) + swa_cache = torch.zeros(2, 128, 1, 512) + backend.bind_kv_caches( + [ + empty_cache, + LayerCache(key=cmp_cache, value=None, swa=swa_cache), + empty_cache, + ] + ) + block_tables = [[], [torch.tensor([[1]], dtype=torch.int32) for _ in range(8)], []] + slot_mappings = [[], [torch.tensor([128], dtype=torch.int32) for _ in range(8)], []] + dsa = SimpleNamespace( + block_tables=block_tables, + slot_mappings=slot_mappings, + actual_seq_lengths_query=torch.tensor([0, 1], dtype=torch.int32), + actual_seq_lengths_kv=torch.tensor([1], dtype=torch.int32), + input_positions=torch.tensor([0], dtype=torch.int64), + cos_table=torch.zeros(1, 64), + sin_table=torch.zeros(1, 64), + c4_cos=torch.zeros(1, 64), + c4_sin=torch.zeros(1, 64), + c128_cos=torch.zeros(1, 64), + c128_sin=torch.zeros(1, 64), + c4_metadata=torch.zeros(1, dtype=torch.int32), + ) + backend._metadata = SimpleNamespace( + dsa_metadata=dsa, + is_prefill=False, + is_chunked_prefill=False, + ) + + with pytest.raises(RuntimeError, match="compressor is required"): + backend.execute( + torch.zeros(1, 8, 512), + torch.zeros(1, 1, 512), + torch.zeros(1, 1, 512), + SimpleNamespace(layer_id=1, attn_sink=None), + ) + + +def test_forward_rope_state_is_owned_by_each_metadata(monkeypatch) -> None: + backend = _make_backend() + monkeypatch.setattr(backend, "_move_metadata_to_device", lambda dsa: None) + monkeypatch.setattr(backend, "_build_precomputed_metadata", lambda dsa, metadata: None) + + def make_metadata(kv_len: int, q_len: int) -> SimpleNamespace: + return SimpleNamespace( + multi_block_tables=[], + kv_seq_lens_host=torch.tensor([kv_len], dtype=torch.int32), + q_seq_lens_host=torch.tensor([q_len], dtype=torch.int32), + is_prefill=q_len > 1, + is_chunked_prefill=False, + dsa_metadata=None, + dsa_positions=None, + dsa_cos_sin=None, + dsa_c4_cos_sin=None, + dsa_c128_cos_sin=None, + dsa_graph_block_table_cols=0, + dsa_graph_mode=False, + ) + + rope_cache = torch.arange(256 * 8, dtype=torch.float32).view(256, 8) + prefill = make_metadata(84, 84) + backend.prepare(prefill) + backend.attach_rope_tables( + torch.arange(84), + rope_cache, + c4_cos_sin=rope_cache, + c128_cos_sin=rope_cache, + metadata=prefill, + ) + + decode = make_metadata(85, 1) + backend.prepare(decode) + backend.attach_rope_tables( + torch.tensor([84]), + rope_cache, + c4_cos_sin=rope_cache, + c128_cos_sin=rope_cache, + metadata=decode, + ) + + backend.prepare_dsa_metadata_for_forward(prefill) + + assert prefill.dsa_metadata.input_positions.numel() == 84 + assert decode.dsa_positions.numel() == 1 + assert prefill.dsa_positions.data_ptr() != decode.dsa_positions.data_ptr() + + +def test_prefill_persists_swa_for_decode_and_omits_ori_kv_cu_seqlens( + monkeypatch, +) -> None: + backend = DsaAttentionBackend( + compress_ratios=[1], + window_size=128, + n_layers=1, + num_heads=8, + attn_head_dim=512, + index_topk=512, + index_n_heads=64, + index_head_dim=128, + rope_head_dim=64, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + swa = torch.zeros(2, 128, 1, 512, dtype=torch.float32) + backend.bind_kv_caches([LayerCache(key=None, value=None, swa=swa)]) + block_table = torch.tensor([[1]], dtype=torch.int32) + layer = SimpleNamespace(layer_id=0, attn_sink=None) + calls: list[dict] = [] + + def fake_sparse_attn(**kwargs): + calls.append(kwargs) + return kwargs["q"].clone(), torch.empty(0) + + monkeypatch.setattr(dsa_attention_module, "_sparse_attn_sharedkv", fake_sparse_attn) + + def prepare_step(kv_len: int, q_len: int, is_prefill: bool): + dsa = backend._builder.build( + multi_block_tables=[block_table], + kv_seq_lens=[kv_len], + q_seq_lens=[q_len], + positions=torch.arange(kv_len - q_len, kv_len, dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=is_prefill, + is_chunked_prefill=False, + ) + dsa.c1_metadata = torch.zeros(1, dtype=torch.int32) + backend._metadata = SimpleNamespace( + dsa_metadata=dsa, + is_prefill=is_prefill, + is_chunked_prefill=False, + ) + return dsa + + prepare_step(kv_len=2, q_len=2, is_prefill=True) + prefill_kv = torch.arange(2 * 512, dtype=torch.float32).view(2, 1, 512) + backend.execute(torch.zeros(2, 8, 512), prefill_kv, prefill_kv, layer) + assert torch.equal(swa[1, :2], prefill_kv) + + prepare_step(kv_len=3, q_len=1, is_prefill=False) + decode_kv = torch.full((1, 1, 512), 7.0) + backend.execute(torch.zeros(1, 8, 512), decode_kv, decode_kv, layer) + assert torch.equal(swa[1, 2], decode_kv[0]) + assert calls[-1]["cu_seqlens_ori_kv"] is None diff --git a/tests/python/test_dsa_metadata.py b/tests/python/test_dsa_metadata.py new file mode 100644 index 0000000000..3831fc0de0 --- /dev/null +++ b/tests/python/test_dsa_metadata.py @@ -0,0 +1,428 @@ +# 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. + +"""Unit tests for the Python DSA metadata builder. + +Validates the faithful port of ``DSAMetadataBuilder`` (core/layers/common/ +dsa_metadata_builder.cpp) against the cache-spec and slot-expansion rules the +C++ implementation enforces. These tests are pure Python: they do not load the +compiled NPU operators, so they run anywhere. +""" + +from __future__ import annotations + +import pytest +import torch + +from xllm.python.attention.dsa_metadata import ( + DSA_CACHE_SLIDING_WINDOW, + DSA_CACHE_TOKEN, + DsaMetadataBuilder, + build_cache_specs, +) + + +def test_build_cache_specs_groups() -> None: + """Group 0 is always SWA; TOKEN groups register in first-seen order.""" + compress_ratios = [0, 0, 4, 128, 4, 128, 4, 0] + caches_info, group_infos = build_cache_specs(compress_ratios, window_size=128, n_layers=8) + + # Three groups: SWA(1,128), TOKEN(4,128), TOKEN(128,128). + assert len(group_infos) == 3 + assert group_infos[0].cache_type == DSA_CACHE_SLIDING_WINDOW + assert group_infos[0].ratio == 1 + assert group_infos[1].cache_type == DSA_CACHE_TOKEN + assert group_infos[1].ratio == 4 + assert group_infos[2].cache_type == DSA_CACHE_TOKEN + assert group_infos[2].ratio == 128 + + +def test_build_cache_specs_per_layer_cache_counts() -> None: + """C1 -> 1 cache, C4 -> 8 caches, C128 -> 4 caches.""" + compress_ratios = [0, 4, 128] + caches_info, _ = build_cache_specs(compress_ratios, 128, 3) + + assert len(caches_info[0]) == 1 # cr=0 -> normalized to 1 + assert len(caches_info[1]) == 8 # cr=4 + assert len(caches_info[2]) == 4 # cr=128 + + +def test_build_cache_specs_does_not_silently_accept_unknown_ratio() -> None: + caches_info, group_infos = build_cache_specs([2], 128, 1) + + assert len(group_infos) == 1 + assert caches_info == [[]] + + +def test_build_cache_specs_real_dsv4_config() -> None: + """The shipped DeepSeek-V4-Flash config produces 2 C1 + 21 C4 + 20 C128. + + config.json has ``compress_ratios`` of length 44 (3 zeros, 21 fours, 20 + one-twenty-eights) but ``num_hidden_layers=43``; only layers 0..42 are + built, so the trailing zero (index 43) is ignored and two C1 layers + remain (indices 0 and 1). + """ + compress_ratios = ( + [0, 0] + [4, 128] * 20 + [4, 0] # layer 42 is C4; index 43 (zero) is ignored. + ) + assert len(compress_ratios) == 44 + caches_info, group_infos = build_cache_specs(compress_ratios, 128, 43) + + assert len(group_infos) == 3 + c1 = sum(1 for layer in caches_info if len(layer) == 1) + c4 = sum(1 for layer in caches_info if len(layer) == 8) + c128 = sum(1 for layer in caches_info if len(layer) == 4) + assert c1 == 2 + assert c4 == 21 + assert c128 == 20 + assert c1 + c4 + c128 == 43 + + +def _make_builder(n_layers: int = 4) -> tuple[DsaMetadataBuilder, list, list]: + compress_ratios = [0, 4, 128, 4] + caches_info, group_infos = build_cache_specs(compress_ratios, 128, n_layers) + return DsaMetadataBuilder(caches_info, group_infos), caches_info, group_infos + + +def test_build_seq_lengths_and_start_pos() -> None: + """start_pos = kv_len - q_len per sequence.""" + builder, _, _ = _make_builder() + # batch=2, decode (q_len=1 each). + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[6, 8], + q_seq_lens=[1, 1], + positions=torch.tensor([5, 7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + assert dsa.seq_lens.tolist() == [6, 8] + assert dsa.seq_lens_q.tolist() == [1, 1] + assert dsa.start_pos.tolist() == [5, 7] + # actual_seq_lengths_query is cumsum(q_lens) with a leading zero. + assert dsa.actual_seq_lengths_query.tolist() == [0, 1, 2] + assert dsa.kv_cu_seq_lens.tolist() == [0, 6, 14] + assert dsa.max_seqlen_q.ndim == 0 + assert dsa.max_seqlen_kv.ndim == 0 + + +def test_build_max_lengths_include_attention_metadata_capacity() -> None: + """C++ takes max(params.meta.max_*, max(host sequence lengths)).""" + builder, _, _ = _make_builder() + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[6, 8], + q_seq_lens=[1, 2], + positions=torch.tensor([5, 6, 7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + max_query_len=16, + max_seq_len=32, + ) + + assert dsa.max_query_len == 16 + assert dsa.max_seq_len == 32 + + +def test_build_token_group_slot_committed_rows() -> None: + """A TOKEN cache commits one row per ratio boundary crossed this step.""" + builder, caches_info, group_infos = _make_builder() + # group 0 = SWA, group 1 = TOKEN(4). Give each a [batch=1, cols=4] table. + swa_bt = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21, 22, 23]], dtype=torch.int32) + # kv_len=8, q_len=1 (decode): prev_ctx_len=7, committed = 8//4 - 7//4 = 2 - 1 = 1. + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[8], + q_seq_lens=[1], + positions=torch.tensor([7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + # Layer 1 (cr=4): cmp cache is caches_info[1][0] -> group 1 (TOKEN4). + cmp_slot = dsa.slot_mappings[1][0] + # One committed row: compressed_idx = prev_committed = 7//4 = 1. + # block_idx = 1 // 128 = 0, block_id = token4_bt[0,0] = 20. + # slot = 20 * 128 + 1 = 2561. + assert cmp_slot.numel() >= 1 + assert cmp_slot[0].item() == 20 * 128 + 1 + + +def test_build_token_group_slot_empty_between_boundaries() -> None: + """Eager decode uses an actual empty tensor when no row is committed.""" + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21, 22, 23]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[129], + q_seq_lens=[1], + positions=torch.tensor([128], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + + assert dsa.slot_mappings[1][0].numel() == 0 + + +def test_build_token_group_slot_commits_at_later_boundary() -> None: + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21, 22, 23]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[132], + q_seq_lens=[1], + positions=torch.tensor([131], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + + assert dsa.slot_mappings[1][0].tolist() == [20 * 128 + 32] + + +def test_build_swa_group_slot_query_tokens_only() -> None: + """A SWA cache writes only the current forward's query token.""" + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21, 22, 23]], dtype=torch.int32) + # kv_len=8, q_len=1, q_start=7, pos=7, block_idx = 7//128 % 4 = 0, + # block_id = swa_bt[0,0] = 10, offset = 7 % 128 = 7 -> slot = 10*128+7 = 1287. + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[8], + q_seq_lens=[1], + positions=torch.tensor([7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + # Layer 0 (cr=1): the single SWA cache -> group 0. + swa_slot = dsa.slot_mappings[0][0] + assert swa_slot[0].item() == 10 * 128 + 7 + + +def test_build_block_tables_shared_within_group() -> None: + """Caches in the same group share the same underlying tensor.""" + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21, 22, 23]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[8], + q_seq_lens=[1], + positions=torch.tensor([7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + # Layer 1 (cr=4): caches 0,1,7 are TOKEN4 (group 1) -> same slot tensor. + assert dsa.slot_mappings[1][0].data_ptr() == dsa.slot_mappings[1][1].data_ptr() + assert dsa.slot_mappings[1][0].data_ptr() == dsa.slot_mappings[1][7].data_ptr() + # Caches 2-6 are SWA (group 0) -> same slot tensor. + assert dsa.slot_mappings[1][2].data_ptr() == dsa.slot_mappings[1][3].data_ptr() + + +def test_build_c4_pad_positions() -> None: + """c4_pad_positions records next_pos-4 when (pos+1) % 4 == 0.""" + builder, _, _ = _make_builder() + # q_len=4, q_start=3 -> positions 3,4,5,6. (pos+1)%4==0 at pos=3 (next=4). + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[7], + q_seq_lens=[4], + positions=torch.tensor([3, 4, 5, 6], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=True, + is_chunked_prefill=False, + ) + # pos=3 -> next_pos=4 -> 4%4==0 -> record 4-4=0. + assert 0 in dsa.c4_pad_positions.tolist() + + +def test_graph_compressed_positions_use_zero_padding() -> None: + """ACL graph position buffers match C++ vector::resize zero fill.""" + builder, _, _ = _make_builder() + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[7], + q_seq_lens=[4], + positions=torch.tensor([3, 4, 5, 6], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=True, + is_chunked_prefill=False, + enable_graph=True, + ) + + assert dsa.c4_pad_positions.tolist() == [0, 0, 0, 0] + assert dsa.c128_pad_positions.tolist() == [0, 0, 0, 0] + + +def test_empty_batch_preserves_cpp_zero_length_buffers() -> None: + builder, _, _ = _make_builder() + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[], + q_seq_lens=[], + positions=torch.empty(0, dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=True, + is_chunked_prefill=False, + ) + + assert dsa.actual_seq_lengths_query.tolist() == [0] + assert dsa.kv_cu_seq_lens.tolist() == [0] + assert dsa.max_seqlen_q.shape == (1,) + assert dsa.max_seqlen_kv.shape == (1,) + assert dsa.max_query_len == 0 + assert dsa.max_seq_len == 0 + assert dsa.c4_pad_positions.numel() == 0 + assert dsa.c128_pad_positions.numel() == 0 + + +def test_build_c128_slot_at_compression_boundary() -> None: + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21]], dtype=torch.int32) + token128_bt = torch.tensor([[30, 31]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt, token128_bt], + kv_seq_lens=[128], + q_seq_lens=[1], + positions=torch.tensor([127], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + + # Layer 2 uses TOKEN(128) for cache 0. The first compressed row is offset 0. + assert dsa.slot_mappings[2][0].tolist() == [30 * 128] + assert dsa.c128_pad_positions.tolist() == [0] + + +def test_multi_batch_slots_are_concatenated_by_sequence() -> None: + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21], [22, 23]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[4, 8], + q_seq_lens=[1, 1], + positions=torch.tensor([3, 7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + + assert dsa.slot_mappings[0][0].tolist() == [10 * 128 + 3, 12 * 128 + 7] + assert dsa.slot_mappings[1][0].tolist() == [20 * 128, 22 * 128 + 1] + + +def test_packed_manager_block_table_is_unpacked() -> None: + builder, _, _ = _make_builder() + packed = torch.tensor([[10, 11], [20, 21], [30, 31]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[packed], + kv_seq_lens=[128], + q_seq_lens=[1], + positions=torch.tensor([127], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + ) + + assert dsa.block_tables[0][0].tolist() == [[10, -1]] + assert dsa.block_tables[1][0].tolist() == [[20, 21]] + assert dsa.block_tables[2][0].tolist() == [[30, 31]] + + +def test_graph_slots_and_block_tables_use_bucket_capacity() -> None: + builder, _, _ = _make_builder() + swa_bt = torch.tensor([[10, 11]], dtype=torch.int32) + token4_bt = torch.tensor([[20, 21]], dtype=torch.int32) + dsa = builder.build( + multi_block_tables=[swa_bt, token4_bt], + kv_seq_lens=[8], + q_seq_lens=[1], + positions=torch.tensor([7, 0, 0, 0], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + enable_graph=True, + graph_block_table_capacity_cols=4, + ) + + assert dsa.slot_mappings[1][0].tolist() == [20 * 128 + 1, -1, -1, -1] + assert dsa.block_tables[1][0].shape == (1, 4) + assert dsa.block_tables[1][0].tolist() == [[20, 21, -1, -1]] + assert dsa.slot_mappings[0][0].tolist() == [10 * 128 + 7, -1, -1, -1] + assert dsa.block_tables[0][0].shape == (1, 4) + + +def test_rope_cache_is_split_into_contiguous_cos_and_sin_tables() -> None: + builder, _, _ = _make_builder() + cos_sin = torch.arange(24, dtype=torch.float32).view(3, 8) + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[3], + q_seq_lens=[3], + positions=torch.arange(3, dtype=torch.int64), + dsa_cos_sin=cos_sin, + is_prefill=True, + is_chunked_prefill=False, + ) + + assert torch.equal(dsa.cos_table, cos_sin[:, :4]) + assert torch.equal(dsa.sin_table, cos_sin[:, 4:]) + assert dsa.cos_table.is_contiguous() + assert dsa.sin_table.is_contiguous() + + +def test_compressed_positions_preserve_position_dtype() -> None: + builder, _, _ = _make_builder() + dsa = builder.build( + multi_block_tables=[], + kv_seq_lens=[4], + q_seq_lens=[4], + positions=torch.arange(4, dtype=torch.int32), + dsa_cos_sin=None, + is_prefill=True, + is_chunked_prefill=False, + ) + + assert dsa.c4_pad_positions.dtype == torch.int32 + assert dsa.c128_pad_positions.dtype == torch.int32 + + +def test_graph_rejects_block_table_larger_than_bucket_capacity() -> None: + builder, _, _ = _make_builder() + block_table = torch.tensor([[10, 11, 12]], dtype=torch.int32) + + with pytest.raises(ValueError, match="exceeds bucket capacity"): + builder.build( + multi_block_tables=[block_table], + kv_seq_lens=[8], + q_seq_lens=[1], + positions=torch.tensor([7], dtype=torch.int64), + dsa_cos_sin=None, + is_prefill=False, + is_chunked_prefill=False, + enable_graph=True, + graph_block_table_capacity_cols=2, + ) diff --git a/tests/python/test_grouped_moe.py b/tests/python/test_grouped_moe.py new file mode 100644 index 0000000000..57cb1751c8 --- /dev/null +++ b/tests/python/test_grouped_moe.py @@ -0,0 +1,123 @@ +# 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. + +"""Contracts for the NPU pre-selected grouped MoE path.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + +_REPO_ROOT = Path(__file__).parents[2] + + +def _load_npu_moe_module(): + path = _REPO_ROOT / "xllm/python/kernels_npu/moe.py" + spec = importlib.util.spec_from_file_location("pr5_npu_moe", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_selected_expert_moe_matches_native_call_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from xllm.python import kernels + + moe = _load_npu_moe_module() + + hidden = torch.empty(3, 16, dtype=torch.bfloat16) + topk_weights = torch.ones(3, 2, dtype=torch.bfloat16) + topk_ids = torch.zeros(3, 2, dtype=torch.int32) + expanded = torch.empty(6, 16, dtype=torch.bfloat16) + row_ids = torch.arange(6, dtype=torch.int32) + group_list = torch.tensor([1, 3, 5, 6], dtype=torch.int64) + quantized = torch.empty(6, 16, dtype=torch.int8) + input_scale = torch.empty(6, dtype=torch.float32) + gemm1 = torch.empty(6, 32, dtype=torch.int32) + activated = torch.empty(6, 16, dtype=torch.int8) + activation_scale = torch.empty(6, dtype=torch.float32) + gemm2 = torch.empty(6, 16, dtype=torch.bfloat16) + calls: list[tuple[str, object]] = [] + + def init_routing(*args, **kwargs): + calls.append(("routing", kwargs)) + return expanded, row_ids, group_list, torch.empty(0) + + def dynamic_quant(value): + assert value is expanded + calls.append(("dynamic_quant", value)) + return quantized, input_scale + + def dequant_swiglu_quant(**kwargs): + calls.append(("dequant_swiglu_quant", kwargs)) + return activated, activation_scale + + gemm_calls: list[dict[str, object]] = [] + + def group_gemm(**kwargs): + gemm_calls.append(kwargs) + return gemm1 if len(gemm_calls) == 1 else gemm2 + + def token_unpermute(**kwargs): + calls.append(("unpermute", kwargs)) + return hidden + + monkeypatch.setattr(moe, "_group_gemm", group_gemm) + monkeypatch.setattr(moe.torch_npu, "npu_moe_init_routing_v2", init_routing) + monkeypatch.setattr(moe.torch_npu, "npu_moe_token_unpermute", token_unpermute) + monkeypatch.setattr(kernels, "dynamic_quant", dynamic_quant, raising=False) + monkeypatch.setattr(kernels, "dequant_swiglu_quant", dequant_swiglu_quant, raising=False) + + result = moe._grouped_moe_with_selected_experts_impl( + hidden, + topk_weights, + topk_ids, + torch.empty(4, 16, 32, dtype=torch.int8), + torch.empty(4, 16, 16, dtype=torch.int8), + torch.empty(4, 32), + torch.empty(4, 16), + num_total_experts=16, + start_expert_id=4, + num_experts_per_rank=4, + swiglu_limit=7.0, + ) + + assert result is hidden + routing = dict(calls)["routing"] + assert isinstance(routing, dict) + assert routing["active_expert_range"] == [4, 8] + assert routing["expert_num"] == 16 + assert routing["quant_mode"] == -1 + + assert len(gemm_calls) == 2 + assert gemm_calls[0]["scale"] is None + assert gemm_calls[0]["per_token_scale"] is None + assert gemm_calls[0]["output_dtype"] == torch.int32 + assert gemm_calls[1]["scale"].dtype == torch.bfloat16 + assert gemm_calls[1]["per_token_scale"] is activation_scale + assert gemm_calls[1]["output_dtype"] == torch.bfloat16 + assert all(call["group_list"] is group_list for call in gemm_calls) + assert all(call["group_list_type"] == 1 for call in gemm_calls) + + dequant = dict(calls)["dequant_swiglu_quant"] + assert isinstance(dequant, dict) + assert dequant["x"] is gemm1 + assert dequant["activation_scale"] is input_scale + assert dequant["group_index"] is group_list + assert dequant["clamp_limit"] == 7.0 diff --git a/tests/python/test_kernels_import.py b/tests/python/test_kernels_import.py index 7fd8a98b5e..d3f256522d 100644 --- a/tests/python/test_kernels_import.py +++ b/tests/python/test_kernels_import.py @@ -31,8 +31,7 @@ _COMMON_SCHEMAS = ( "rms_norm(Tensor input, Tensor weight, float eps) -> Tensor", - "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor weight, " - "float eps) -> (Tensor, Tensor)", + "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor weight, float eps) -> (Tensor, Tensor)", "silu_and_mul(Tensor input) -> Tensor", "reshape_paged_cache(Tensor slot_mapping, Tensor keys, Tensor values, " "Tensor(a!) key_cache, Tensor(b!) value_cache) -> Tensor", @@ -55,10 +54,12 @@ "quant_matmul(Tensor x1, Tensor x2, bool transpose2, Tensor scale, " "Tensor? offset, Tensor? pertoken_scale, Tensor? bias, ScalarType? " "output_dtype) -> Tensor", - "quantize_per_tensor(Tensor self, Tensor scales, Tensor zero_points, " - "ScalarType dtype, int axis) -> Tensor", + "quantize_per_tensor(Tensor self, Tensor scales, Tensor zero_points, ScalarType dtype, int axis) -> Tensor", "dynamic_quant(Tensor input, Tensor? smooth_scales, Tensor? group_index, " "ScalarType? dst_type) -> (Tensor, Tensor?)", + "group_gemm(Tensor x, Tensor weight, Tensor? scale, Tensor? " + "per_token_scale, Tensor group_list, int split_item, int group_type, int " + "group_list_type, ScalarType? output_dtype) -> Tensor", "lightning_indexer(Tensor query, Tensor key, Tensor weights, Tensor? " "query_seq_lengths, Tensor? key_seq_lengths, Tensor? block_table, str " "layout_query, str layout_key, int selected_count, int sparse_mode, int " @@ -79,6 +80,17 @@ "Tensor? actual_seq_lengths_kv, Tensor? query_rope, Tensor? key_rope, " "float scale_value, int sparse_block_size, str layout_query, str layout_kv, " "int sparse_mode, Tensor(a!) output) -> Tensor", + "rms_norm_dynamic_quant(Tensor input, Tensor weight, float eps) -> (Tensor, Tensor)", + "npu_inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, str rotary_mode, int[] partial_slice) -> ()", + "moe_gating_top_k_hash(Tensor x, int k, Tensor? bias, Tensor? input_ids, Tensor? tid2eid, int k_group, int group_count, float routed_scaling_factor, float eps, int group_select_mode, int renorm, int norm_type, bool out_flag) -> (Tensor, Tensor, Tensor)", + "dequant_swiglu_quant(Tensor x, Tensor? weight_scale, Tensor? activation_scale, Tensor? bias, Tensor? quant_scale, Tensor? quant_offset, Tensor? group_index, bool activate_left, int quant_mode, int swiglu_mode, float clamp_limit, float glu_alpha, float glu_bias) -> (Tensor, Tensor)", + "hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, int hc_mult, int hc_sinkhorn_iters, float norm_eps, float hc_eps) -> (Tensor, Tensor, Tensor)", + "hc_post(Tensor x, Tensor residual, Tensor post, Tensor comb) -> Tensor", + "compressor(Tensor x, Tensor wkv, Tensor wgate, Tensor(a!) kv_state, Tensor(b!) score_state, Tensor ape, Tensor norm_weight, Tensor rope_sin, Tensor rope_cos, Tensor? kv_block_table, Tensor? score_block_table, Tensor? cu_seqlens, Tensor? seqused, Tensor? start_pos, int rope_head_dim, int cmp_ratio, int coff, float norm_eps, int rotary_mode, bool enable_grad) -> (Tensor, Tensor, Tensor, Tensor, Tensor)", + "sparse_attn_sharedkv(Tensor q, Tensor? ori_kv, Tensor? cmp_kv, Tensor? ori_sparse_indices, Tensor? cmp_sparse_indices, Tensor? ori_block_table, Tensor? cmp_block_table, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, Tensor? sinks, Tensor? metadata, float softmax_scale, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool return_softmax_lse) -> (Tensor, Tensor)", + "sparse_attn_sharedkv_metadata(int num_heads_q, int num_heads_kv, int head_dim, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, int batch_size, int max_seqlen_q, int max_seqlen_kv, int ori_topk, int cmp_topk, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool has_ori_kv, bool has_cmp_kv) -> Tensor", + "quant_lightning_indexer(Tensor query, Tensor key, Tensor weights, Tensor query_dequant_scale, Tensor key_dequant_scale, int query_quant_mode, int key_quant_mode, Tensor? actual_seq_lengths_query, Tensor? actual_seq_lengths_key, Tensor? block_table, Tensor? metadata, str layout_query, str layout_key, int sparse_count, int sparse_mode, int pre_tokens, int next_tokens, int cmp_ratio, bool return_value) -> (Tensor, Tensor)", + "quant_lightning_indexer_metadata(int num_heads_q, int num_heads_k, int head_dim, int query_quant_mode, int key_quant_mode, Tensor? actual_seq_lengths_query, Tensor? actual_seq_lengths_key, int batch_size, int max_seqlen_q, int max_seqlen_k, str layout_query, str layout_key, int sparse_count, int sparse_mode, int pre_tokens, int next_tokens, int cmp_ratio, str device) -> Tensor", ) _PLATFORM_REQUIRED = pytest.mark.skipif( @@ -128,9 +140,7 @@ def _run_isolated_python( definitions.append(_package_stub(kernel_package)) env = os.environ.copy() - env["PYTHONPATH"] = os.pathsep.join( - value for value in (str(_REPO_ROOT), env.get("PYTHONPATH", "")) if value - ) + env["PYTHONPATH"] = os.pathsep.join(value for value in (str(_REPO_ROOT), env.get("PYTHONPATH", "")) if value) result = subprocess.run( [sys.executable, "-c", "\n".join((*definitions, textwrap.dedent(script)))], cwd=_REPO_ROOT, @@ -143,13 +153,9 @@ def _run_isolated_python( def test_platform_queries_are_no_argument(monkeypatch: pytest.MonkeyPatch) -> None: current_platform = platform.current_platform - monkeypatch.setattr( - type(current_platform), "enum", classmethod(lambda cls: platform.PlatformEnum.NPU) - ) + monkeypatch.setattr(type(current_platform), "enum", classmethod(lambda cls: platform.PlatformEnum.NPU)) assert current_platform.is_npu() and not current_platform.is_cuda() - monkeypatch.setattr( - type(current_platform), "enum", classmethod(lambda cls: platform.PlatformEnum.CUDA) - ) + monkeypatch.setattr(type(current_platform), "enum", classmethod(lambda cls: platform.PlatformEnum.CUDA)) assert current_platform.is_cuda() and not current_platform.is_npu() @@ -201,11 +207,12 @@ def test_registry_does_not_preload_model_modules() -> None: def test_npu_fake_tensor_and_mutation_contracts() -> None: - """Quantization and sparse attention shapes traced without an NPU.""" + """NPU wrapper shape and mutation contracts traced without an NPU.""" _run_isolated_python( """ import xllm.python.kernels_npu._custom_op # noqa: F401 - from xllm.python.kernels_npu import quantization, sparse_attention + from xllm.python.kernels_npu import dsa, normalization, quantization + from xllm.python.kernels_npu import rotary_embedding, sparse_attention mode = torch._subclasses.fake_tensor.FakeTensorMode() with mode: @@ -221,6 +228,37 @@ def test_npu_fake_tensor_and_mutation_contracts() -> None: assert quantized.shape == (2, 2) and quantized.dtype == torch.int32 assert scale.shape == (2,) and scale.dtype == torch.float32 + grouped = torch.ops.xllm_ops.group_gemm( + torch.empty(6, 16, dtype=torch.int8), + torch.empty(4, 16, 32, dtype=torch.int8), + None, + None, + torch.empty(4, dtype=torch.int64), + 2, + 0, + 1, + torch.int32, + ) + assert grouped.shape == (6, 32) + assert grouped.dtype == torch.int32 + + from xllm.python.kernels_npu import moe + + selected_moe = moe.grouped_moe_with_selected_experts( + torch.empty(3, 16, dtype=torch.bfloat16), + torch.empty(3, 2, dtype=torch.bfloat16), + torch.empty(3, 2, dtype=torch.int32), + torch.empty(4, 16, 32, dtype=torch.int8), + torch.empty(4, 16, 16, dtype=torch.int8), + torch.empty(4, 32), + torch.empty(4, 16), + num_total_experts=16, + start_expert_id=4, + num_experts_per_rank=4, + ) + assert selected_moe.shape == (3, 16) + assert selected_moe.dtype == torch.bfloat16 + query = torch.empty(8, 4, 16) key = torch.empty(8, 2, 16) indices = sparse_attention.lightning_indexer( @@ -237,6 +275,107 @@ def test_npu_fake_tensor_and_mutation_contracts() -> None: torch.empty(2, 2, 16), ) is None + normed, norm_scale = normalization.rms_norm_dynamic_quant( + torch.empty(8, 16), torch.empty(16), 1e-6 + ) + assert normed.shape == (8, 16) and normed.dtype == torch.int8 + assert norm_scale.shape == (8,) and norm_scale.dtype == torch.float32 + + rotary_input = torch.empty(8, 2, 128) + rotary_ptr = rotary_input.data_ptr() + assert rotary_embedding.npu_inplace_partial_rotary_mul( + rotary_input, torch.empty(8, 64), torch.empty(8, 64), 64, 64 + ).data_ptr() == rotary_ptr + + compressor_out = dsa.compressor( + torch.empty(8, 16), + torch.empty(8, 16), + torch.empty(8, 16), + torch.empty(1, 128, 8), + torch.empty(1, 128, 8), + torch.empty(4, 8), + torch.empty(8), + torch.empty(2, 4), + torch.empty(2, 4), + None, None, None, None, None, + 4, 4, 1, 1e-6, 1, False, + ) + assert compressor_out[0].shape == (2, 8) + assert all(tensor.numel() == 0 for tensor in compressor_out[1:]) + + seq_lens = torch.empty(1, dtype=torch.int32) + sparse_metadata = dsa.sparse_attn_sharedkv_metadata( + 64, 1, 512, None, None, None, seq_lens, seq_lens, + 1, 4, 16, 0, 0, 1, 4, 3, 127, 0, + "BSND", "PA_ND", True, False, + ) + assert sparse_metadata.shape == (1024,) + assert sparse_metadata.dtype == torch.int32 + + dsa_query = torch.empty(1, 4, 64, 512) + sparse_out, sparse_lse = dsa.sparse_attn_sharedkv( + dsa_query, None, None, None, None, None, None, + None, None, None, None, None, None, sparse_metadata, + 1.0, 1, 4, 3, 127, 0, "BSND", "PA_ND", False, + ) + assert sparse_out.shape == dsa_query.shape + assert sparse_lse.numel() == 0 + + qli_metadata = dsa.quant_lightning_indexer_metadata( + 64, 1, 128, 0, 0, seq_lens, seq_lens, + 1, 8, 8, "TND", "PA_BSND", 512, 3, + 2**63 - 1, 2**63 - 1, 4, "cpu", + ) + assert qli_metadata.shape == (1024,) + assert qli_metadata.dtype == torch.int32 + + qli_indices, qli_values = dsa.quant_lightning_indexer( + torch.empty(8, 64, 128, dtype=torch.int8), + torch.empty(1, 128, 1, 128, dtype=torch.int8), + torch.empty(8, 64), + torch.empty(8, 64), + torch.empty(1, 128, 1), + 0, 0, seq_lens, seq_lens, torch.empty(1, 1), qli_metadata, + "TND", "PA_BSND", 512, 3, + 2**63 - 1, 2**63 - 1, 4, False, + ) + assert qli_indices.shape == (8, 1, 512) + assert qli_indices.dtype == torch.int32 + assert qli_values.numel() == 0 + + hc_input = torch.empty(8, 4, 16) + hc_attn, hc_post, hc_comb = dsa.hc_pre( + hc_input, + torch.empty(24, 64), + torch.empty(3), + torch.empty(24), + 4, 20, 1e-6, 1e-6, + ) + assert hc_attn.shape == (8, 16) + assert hc_post.shape == (8, 4) + assert hc_comb.shape == (8, 4, 4) + assert dsa.hc_post( + hc_attn, hc_input, hc_post, hc_comb + ).shape == hc_input.shape + + gate_weights, expert_idx, gate_out = dsa.moe_gating_top_k_hash( + torch.empty(8, 256), 6, None, None, None, + 1, 1, 1.0, 1e-20, 1, 0, 2, False, + ) + assert gate_weights.shape == (8, 6) + assert expert_idx.shape == (8, 6) + assert expert_idx.dtype == torch.int32 + assert gate_out.shape == (8, 256) + assert gate_out.dtype == torch.float32 + + swiglu_out, swiglu_scale = dsa.dequant_swiglu_quant( + torch.empty(8, 32, dtype=torch.int32), None, None + ) + assert swiglu_out.shape == (8, 16) + assert swiglu_out.dtype == torch.int8 + assert swiglu_scale.shape == (8,) + assert swiglu_scale.dtype == torch.float32 + try: quantization.dynamic_quant( torch.empty(2, 15), dst_type=torch.quint4x2 diff --git a/tests/python/test_model_executor.py b/tests/python/test_model_executor.py index e07abd20fd..362ee5d7ea 100644 --- a/tests/python/test_model_executor.py +++ b/tests/python/test_model_executor.py @@ -37,6 +37,7 @@ AttentionBackend, AttentionMetadata, LayerCache, + normalize_layer_caches, ) from xllm.python.layers.attention import Attention # noqa: E402 from xllm.python.model_executor.executor import ( # noqa: E402 @@ -44,14 +45,13 @@ _create_attention_backend, _resolve_graph_backend, ) +from xllm.python.model_executor.runners.decode_acl_graph import ( # noqa: E402 + DecodeAclGraphRunner, +) from xllm.python.model_executor.runners.decode_cuda_graph import ( # noqa: E402 DecodeCudaGraphRunner, _decode_graph_buckets, ) -from xllm.python.model_executor.runners.decode_acl_graph import ( # noqa: E402 - DecodeAclGraphRunner, -) - # --------------------------------------------------------------------------- # Helpers @@ -91,7 +91,12 @@ def page_size(self) -> int: def _make_attention_layer( - num_heads=8, num_kv_heads=2, head_dim=64, scale=0.125, sliding_window=0, layer_id=0, + num_heads=8, + num_kv_heads=2, + head_dim=64, + scale=0.125, + sliding_window=0, + layer_id=0, ) -> Attention: return Attention( num_heads=num_heads, @@ -109,9 +114,7 @@ class _FakeModel(nn.Module): def __init__(self, num_layers: int = 2, device: str = "cpu", **attn_kwargs): super().__init__() self.model = nn.Linear(1, 1) # execution_model placeholder - self.layers = nn.ModuleList( - [_make_attention_layer(layer_id=i, **attn_kwargs) for i in range(num_layers)] - ) + self.layers = nn.ModuleList([_make_attention_layer(layer_id=i, **attn_kwargs) for i in range(num_layers)]) self._param = nn.Parameter(torch.zeros(1, device=device)) def forward(self, input_ids, positions): @@ -159,6 +162,32 @@ def test_enable_graph_selects_aclgraph_on_npu(self, _mock_is_npu): class TestCreateAttentionBackend: + @patch( + "xllm.python.model_executor.executor.current_platform.is_npu", + return_value=True, + ) + def test_deepseek_v4_creates_dsa_backend(self, _mock_is_npu): + attn = _make_attention_layer(head_dim=512) + module = types.ModuleType("xllm.python.attention.dsa_attention") + module.DsaAttentionBackend = StubAttentionBackend + config = { + "model_type": "deepseek_v4", + "compress_ratios": [1, 4, 128], + "num_hidden_layers": 3, + "window_size": 128, + "index_topk": 512, + "index_n_heads": 64, + "index_head_dim": 128, + "qk_rope_head_dim": 64, + } + with patch.dict(sys.modules, {module.__name__: module}): + backend = _create_attention_backend(attn, torch.device("npu"), torch.bfloat16, config) + + assert isinstance(backend, StubAttentionBackend) + assert backend.init_kwargs["attn_head_dim"] == 512 + assert backend.init_kwargs["n_layers"] == 3 + assert backend.init_kwargs["compress_ratios"] == [1, 4, 128] + @patch( "xllm.python.model_executor.executor.current_platform.is_npu", return_value=True, @@ -169,9 +198,7 @@ class TestCreateAttentionBackend: ) def test_npu_device_creates_npu_backend(self, _mock_is_npu): attn = _make_attention_layer() - backend = _create_attention_backend( - attn, torch.device("npu"), torch.float16 - ) + backend = _create_attention_backend(attn, torch.device("npu"), torch.float16) assert isinstance(backend, StubAttentionBackend) assert backend.init_kwargs["num_heads"] == 8 assert backend.init_kwargs["num_kv_heads"] == 2 @@ -185,16 +212,12 @@ def test_npu_device_creates_npu_backend(self, _mock_is_npu): "xllm.python.model_executor.executor.current_platform.is_cuda", return_value=True, ) - def test_cuda_device_creates_flashinfer_backend( - self, _mock_is_cuda, _mock_is_npu - ): + def test_cuda_device_creates_flashinfer_backend(self, _mock_is_cuda, _mock_is_npu): attn = _make_attention_layer() module = types.ModuleType("xllm.python.attention.flashinfer") module.FlashInferBackend = StubAttentionBackend with patch.dict(sys.modules, {module.__name__: module}): - backend = _create_attention_backend( - attn, torch.device("cuda"), torch.float16 - ) + backend = _create_attention_backend(attn, torch.device("cuda"), torch.float16) assert isinstance(backend, StubAttentionBackend) @@ -242,20 +265,13 @@ def test_heterogeneous_attention_raises(self, _mock_backend): def test_graph_backend_off_variants(self, _mock_backend): for off_value in ("off", "", "none", "0"): model = _FakeModel(num_layers=1) - executor = ModelExecutor( - model, {"python_graph_backend": off_value}, max_seqs_per_batch=4 - ) + executor = ModelExecutor(model, {"python_graph_backend": off_value}, max_seqs_per_batch=4) assert executor.decode_graph_runner is None assert executor.inductor_runner is None - @patch( - "xllm.python.model_executor.runners.decode_cuda_graph." - "DecodeCudaGraphRunner" - ) + @patch("xllm.python.model_executor.runners.decode_cuda_graph.DecodeCudaGraphRunner") @patch("xllm.python.model_executor.executor._create_attention_backend") - def test_data_parallel_cuda_graph_is_supported( - self, mock_create, mock_graph_runner - ): + def test_data_parallel_cuda_graph_is_supported(self, mock_create, mock_graph_runner): mock_create.return_value = StubAttentionBackend() model = _FakeModel(num_layers=1) @@ -334,12 +350,8 @@ def test_single_rank_graph_key_reuses_padded_bucket(self): runner.dp_size = 1 runner.dp_rank = 0 - first = runner._graph_key( - torch.zeros(3, dtype=torch.int32), self._metadata([3]) - ) - second = runner._graph_key( - torch.zeros(4, dtype=torch.int32), self._metadata([4]) - ) + first = runner._graph_key(torch.zeros(3, dtype=torch.int32), self._metadata([3])) + second = runner._graph_key(torch.zeros(4, dtype=torch.int32), self._metadata([4])) assert first == (4, (4,)) assert second == first @@ -370,9 +382,7 @@ def test_can_execute_requires_warmed_graph(self): assert runner.can_execute(input_ids, metadata) @pytest.mark.parametrize("token_counts", ([3], [3, -1], [3, 2])) - def test_graph_key_rejects_invalid_data_parallel_metadata( - self, token_counts - ): + def test_graph_key_rejects_invalid_data_parallel_metadata(self, token_counts): runner = self._runner(dp_rank=1) input_ids = torch.zeros(1, dtype=torch.int32) @@ -399,46 +409,28 @@ def _runner() -> DecodeAclGraphRunner: def _metadata() -> SimpleNamespace: return SimpleNamespace( slot_mapping=torch.arange(4, dtype=torch.int32), - paged_kv_indptr=torch.tensor( - [0, 1, 2, 4, 6], dtype=torch.int32 - ), - paged_kv_indices=torch.tensor( - [10, 10, 20, 21, 20, 21], dtype=torch.int32 - ), - paged_kv_last_page_len=torch.tensor( - [3, 4, 3, 4], dtype=torch.int32 - ), + paged_kv_indptr=torch.tensor([0, 1, 2, 4, 6], dtype=torch.int32), + paged_kv_indices=torch.tensor([10, 10, 20, 21, 20, 21], dtype=torch.int32), + paged_kv_last_page_len=torch.tensor([3, 4, 3, 4], dtype=torch.int32), q_cu_seq_lens=torch.tensor([0, 2, 4], dtype=torch.int32), kv_cu_seq_lens=torch.tensor([0, 4, 12], dtype=torch.int32), kv_seq_lens_host=torch.tensor([4, 8], dtype=torch.int32), kv_seq_lens_host_values=[4, 8], - block_table=torch.tensor( - [[10, 11], [20, 21]], dtype=torch.int32 - ), + block_table=torch.tensor([[10, 11], [20, 21]], dtype=torch.int32), kv_seq_lens=torch.tensor([4, 8], dtype=torch.int32), q_seq_lens=torch.tensor([2, 2], dtype=torch.int32), expanded_decode_metadata=SimpleNamespace( enabled=True, - kv_seq_lens=torch.tensor( - [3, 4, 7, 8], dtype=torch.int32 - ), + kv_seq_lens=torch.tensor([3, 4, 7, 8], dtype=torch.int32), block_table=torch.tensor( [[10, 11], [10, 11], [20, 21], [20, 21]], dtype=torch.int32, ), - paged_kv_indptr=torch.tensor( - [0, 1, 2, 4, 6], dtype=torch.int32 - ), - paged_kv_indices=torch.tensor( - [10, 10, 20, 21, 20, 21], dtype=torch.int32 - ), - paged_kv_last_page_len=torch.tensor( - [3, 4, 3, 4], dtype=torch.int32 - ), + paged_kv_indptr=torch.tensor([0, 1, 2, 4, 6], dtype=torch.int32), + paged_kv_indices=torch.tensor([10, 10, 20, 21, 20, 21], dtype=torch.int32), + paged_kv_last_page_len=torch.tensor([3, 4, 3, 4], dtype=torch.int32), paged_attention_tiling_data=None, - kv_seq_lens_host=torch.tensor( - [3, 4, 7, 8], dtype=torch.int32 - ), + kv_seq_lens_host=torch.tensor([3, 4, 7, 8], dtype=torch.int32), kv_seq_lens_host_values=[3, 4, 7, 8], ), is_prefill=False, @@ -604,6 +596,52 @@ def test_token_layout_mismatch_fails( # --------------------------------------------------------------------------- +class TestNormalizeLayerCaches: + def test_legacy_five_slot_cache_keeps_generic_layout(self): + tensors = tuple(torch.full((1,), value) for value in range(1, 6)) + + cache = normalize_layer_caches([tensors])[0] + + assert cache.key is tensors[0] + assert cache.value is tensors[1] + assert cache.index is tensors[2] + assert cache.conv is tensors[3] + assert cache.ssm is tensors[4] + assert cache.swa is None + assert cache.compress_kv_state is None + assert cache.compress_score_state is None + assert cache.compress_index_kv_state is None + assert cache.compress_index_score_state is None + assert cache.indexer_scale is None + + def test_deepseek_v4_eleven_slot_cache_maps_all_slots(self): + tensors = tuple(torch.full((1,), value) for value in range(1, 12)) + + cache = normalize_layer_caches([tensors])[0] + + assert ( + cache.key, + cache.value, + cache.index, + cache.conv, + cache.ssm, + cache.swa, + cache.compress_kv_state, + cache.compress_score_state, + cache.compress_index_kv_state, + cache.compress_index_score_state, + cache.indexer_scale, + ) == tensors + + def test_empty_deepseek_v4_slots_are_normalized_to_none(self): + cache = normalize_layer_caches([(torch.ones(1), torch.ones(1), *(torch.empty(0),) * 9)])[0] + + assert cache.key is not None + assert cache.value is not None + assert cache.index is None + assert cache.indexer_scale is None + + class TestBindKvCaches: @patch( "xllm.python.model_executor.executor._create_attention_backend", diff --git a/xllm/core/kernels/npu/CMakeLists.txt b/xllm/core/kernels/npu/CMakeLists.txt index c5905568d6..bf8b722fba 100644 --- a/xllm/core/kernels/npu/CMakeLists.txt +++ b/xllm/core/kernels/npu/CMakeLists.txt @@ -37,6 +37,7 @@ cc_library( w4a8_dynamic_moe_preprocess.cpp rec_constrained_topk.cpp npu_ops_library.cpp + group_gemm_wrapper.cpp DEPS :torch_npu_kernels :tilelang_kernels diff --git a/xllm/core/kernels/npu/group_gemm_wrapper.cpp b/xllm/core/kernels/npu/group_gemm_wrapper.cpp new file mode 100644 index 0000000000..29ed7ea56e --- /dev/null +++ b/xllm/core/kernels/npu/group_gemm_wrapper.cpp @@ -0,0 +1,65 @@ +/* Copyright 2025-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/jd-opensource/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. +==============================================================================*/ + +#include "npu_ops_api.h" + +namespace xllm::kernel::npu { + +torch::Tensor group_gemm(const torch::Tensor& x, + const torch::Tensor& weight, + const std::optional& scale, + const std::optional& per_token_scale, + const torch::Tensor& group_list, + int64_t split_item, + int64_t group_type, + int64_t group_list_type, + std::optional output_dtype) { + std::vector x_list = {x}; + std::vector weight_list = {weight}; + std::vector scale_storage; + std::vector per_token_scale_storage; + std::optional scale_list = std::nullopt; + if (scale.has_value()) { + scale_storage.push_back(scale.value()); + scale_list = torch::TensorList(scale_storage); + } + std::optional per_token_scale_list = std::nullopt; + if (per_token_scale.has_value()) { + per_token_scale_storage.push_back(per_token_scale.value()); + per_token_scale_list = torch::TensorList(per_token_scale_storage); + } + auto outputs = + apply_npu_grouped_matmul(torch::TensorList(x_list), + torch::TensorList(weight_list), + /*bias=*/std::nullopt, + scale_list, + /*offset=*/std::nullopt, + /*antiquant_scale=*/std::nullopt, + /*antiquant_offset=*/std::nullopt, + per_token_scale_list, + group_list, + /*activation_input=*/std::nullopt, + /*activation_quant_scale=*/std::nullopt, + /*activation_quant_offset=*/std::nullopt, + split_item, + group_type, + group_list_type, + /*act_type=*/std::nullopt, + /*tuning_config=*/c10::nullopt, + output_dtype); + return outputs.back(); +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index 8b5889ad7f..fb85428618 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -446,4 +446,15 @@ std::tuple apply_npu_mega_moe( int64_t dispatch_quant_out_dtype = 0, int64_t topo_type = 0, int64_t rank_num_per_server = 2); + +torch::Tensor group_gemm(const torch::Tensor& x, + const torch::Tensor& weight, + const std::optional& scale, + const std::optional& per_token_scale, + const torch::Tensor& group_list, + int64_t split_item, + int64_t group_type, + int64_t group_list_type, + std::optional output_dtype); + } // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index 1c9ae56305..110d975a58 100644 --- a/xllm/core/kernels/npu/npu_ops_library.cpp +++ b/xllm/core/kernels/npu/npu_ops_library.cpp @@ -283,6 +283,19 @@ TORCH_LIBRARY(xllm_ops, m) { "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor " "weight, " "float eps) -> (Tensor, Tensor)"); + // Fused RMSNorm + dynamic per-token int8 quant (W8A8 query preprocess). + // Returns (qr_int8, qr_pertoken_scale) matching C++ rms_norm_dynamic_quant + // (npu_ops_api.h:122), used by the DSV4 indexer build_query path. + m.def( + "rms_norm_dynamic_quant(Tensor input, Tensor weight, float eps) -> " + "(Tensor, Tensor)"); + // In-place partial rotary embedding (interleaved). x is 4D [B,N,S,D], r1/r2 + // are cos/sin [B,1,1,rope_head_dim]; partial_slice=[rope_start, + // rope_head_dim]. Mirrors C++ apply_partial_rope + // (deepseek_sparse_attention.cpp:151) used by the DSV4 indexer build_query. + m.def( + "npu_inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, " + "str rotary_mode, int[] partial_slice) -> ()"); m.def("silu_and_mul(Tensor input) -> Tensor"); m.def( "fused_qk_norm_rope(Tensor(a!) qkv, int num_heads_q, int num_heads_k, " @@ -314,6 +327,10 @@ TORCH_LIBRARY(xllm_ops, m) { m.def( "dynamic_quant(Tensor input, Tensor? smooth_scales, Tensor? group_index, " "ScalarType? dst_type) -> (Tensor, Tensor?)"); + m.def( + "group_gemm(Tensor x, Tensor weight, Tensor? scale, Tensor? " + "per_token_scale, Tensor group_list, int split_item, int group_type, int " + "group_list_type, ScalarType? output_dtype) -> Tensor"); m.def( "lightning_indexer(Tensor query, Tensor key, Tensor weights, " "Tensor? query_seq_lengths, Tensor? key_seq_lengths, Tensor? " @@ -349,11 +366,84 @@ TORCH_LIBRARY(xllm_ops, m) { "shard_valid_mask, Tensor restore_index, Tensor query_index, Tensor " "kv_gather_index, int[] q_cu_seqlens, int[] kv_cu_seqlens, int " "total_local)"); + // ---- DeepSeek-V4 DSA kernels ---- + // MoE hash routing gate (returns routed output, expert_idx, token_unpermute). + m.def( + "moe_gating_top_k_hash(Tensor x, int k, Tensor? bias, Tensor? input_ids, " + "Tensor? tid2eid, int k_group, int group_count, float " + "routed_scaling_factor, " + "float eps, int group_select_mode, int renorm, int norm_type, bool " + "out_flag) -> (Tensor, Tensor, Tensor)"); + // Dequant + SwiGLU + quant (fused, replaces manual dequant loop). + m.def( + "dequant_swiglu_quant(Tensor x, Tensor? weight_scale, Tensor? " + "activation_scale, Tensor? bias, Tensor? quant_scale, Tensor? " + "quant_offset, Tensor? group_index, bool activate_left, int quant_mode, " + "int swiglu_mode, float clamp_limit, float glu_alpha, float glu_bias) " + "-> (Tensor, Tensor)"); + // HyperConnection pre/post (hc_pre returns attn_input, post, comb). + m.def( + "hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, " + "int hc_mult, int hc_sinkhorn_iters, float norm_eps, float hc_eps) " + "-> (Tensor, Tensor, Tensor)"); + m.def( + "hc_post(Tensor x, Tensor residual, Tensor post, Tensor comb) -> " + "Tensor"); + // Compressor: NSA-style KV pooling. kv_state/score_state are in-place (Ref). + // Returns (cmp_kv, wkv_proj, softmax_res, norm_x, norm_rstd). + m.def( + "compressor(Tensor x, Tensor wkv, Tensor wgate, Tensor(a!) kv_state, " + "Tensor(b!) score_state, Tensor ape, Tensor norm_weight, Tensor " + "rope_sin, Tensor rope_cos, Tensor? kv_block_table, Tensor? " + "score_block_table, Tensor? cu_seqlens, Tensor? seqused, Tensor? " + "start_pos, int rope_head_dim, int cmp_ratio, int coff, float " + "norm_eps, int rotary_mode, bool enable_grad) -> (Tensor, Tensor, " + "Tensor, Tensor, Tensor)"); + // Two-stage sparse attention over original + compressed KV. + m.def( + "sparse_attn_sharedkv(Tensor q, Tensor? ori_kv, Tensor? cmp_kv, " + "Tensor? ori_sparse_indices, Tensor? cmp_sparse_indices, Tensor? " + "ori_block_table, Tensor? cmp_block_table, Tensor? cu_seqlens_q, " + "Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? " + "seqused_q, Tensor? seqused_kv, Tensor? sinks, Tensor? metadata, " + "float softmax_scale, int cmp_ratio, int ori_mask_mode, int " + "cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, " + "str layout_kv, bool return_softmax_lse) -> (Tensor, Tensor)"); + // AICPU tiling metadata builder for sparse_attn_sharedkv. + m.def( + "sparse_attn_sharedkv_metadata(int num_heads_q, int num_heads_kv, int " + "head_dim, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? " + "cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, int " + "batch_size, int max_seqlen_q, int max_seqlen_kv, int ori_topk, int " + "cmp_topk, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int " + "ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool " + "has_ori_kv, bool has_cmp_kv) -> Tensor"); + // Quantized lightning indexer: int8 q/k top-k selection with cmp_ratio. + m.def( + "quant_lightning_indexer(Tensor query, Tensor key, Tensor weights, " + "Tensor query_dequant_scale, Tensor key_dequant_scale, int " + "query_quant_mode, int key_quant_mode, Tensor? actual_seq_lengths_query, " + "Tensor? actual_seq_lengths_key, Tensor? block_table, Tensor? metadata, " + "str layout_query, str layout_key, int sparse_count, int sparse_mode, " + "int pre_tokens, int next_tokens, int cmp_ratio, bool return_value) -> " + "(Tensor, Tensor)"); + // AICPU tiling metadata builder for quant_lightning_indexer. + m.def( + "quant_lightning_indexer_metadata(int num_heads_q, int num_heads_k, int " + "head_dim, int query_quant_mode, int key_quant_mode, Tensor? " + "actual_seq_lengths_query, Tensor? actual_seq_lengths_key, int " + "batch_size, int max_seqlen_q, int max_seqlen_k, str layout_query, str " + "layout_key, int sparse_count, int sparse_mode, int pre_tokens, int " + "next_tokens, int cmp_ratio, str device) -> Tensor"); } TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("rms_norm", TORCH_FN(xllm::rms_norm_npu)); m.impl("fused_add_rms_norm", TORCH_FN(xllm::fused_add_rms_norm_npu)); + m.impl("rms_norm_dynamic_quant", + TORCH_FN(xllm::kernel::npu::rms_norm_dynamic_quant)); + m.impl("npu_inplace_partial_rotary_mul", + TORCH_FN(xllm::kernel::npu::npu_inplace_partial_rotary_mul)); m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu)); m.impl("reshape_paged_cache", TORCH_FN(xllm::reshape_paged_cache_npu)); m.impl("apply_rotary_embedding", TORCH_FN(xllm::apply_rotary_embedding_npu)); @@ -363,6 +453,7 @@ TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("quantize_per_tensor", TORCH_FN(xllm::kernel::npu::quantize_per_tensor)); m.impl("dynamic_quant", TORCH_FN(xllm::kernel::npu::dynamic_quant)); + m.impl("group_gemm", TORCH_FN(xllm::kernel::npu::group_gemm)); m.impl("lightning_indexer", TORCH_FN(xllm::kernel::npu::lightning_indexer)); m.impl("lightning_indexer_out", TORCH_FN(xllm::kernel::npu::lightning_indexer_out)); @@ -379,4 +470,20 @@ TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { // graph capture), so it needs no fake/meta registration. TORCH_LIBRARY_IMPL(xllm_ops, CompositeExplicitAutograd, m) { m.impl("build_cp_context", TORCH_FN(xllm::build_cp_context_npu)); + // ---- DeepSeek-V4 DSA kernels ---- + m.impl("moe_gating_top_k_hash", + TORCH_FN(xllm::kernel::npu::moe_gating_top_k_hash)); + m.impl("dequant_swiglu_quant", + TORCH_FN(xllm::kernel::npu::dequant_swiglu_quant)); + m.impl("hc_pre", TORCH_FN(xllm::kernel::npu::hc_pre)); + m.impl("hc_post", TORCH_FN(xllm::kernel::npu::hc_post)); + m.impl("compressor", TORCH_FN(xllm::kernel::npu::compressor)); + m.impl("sparse_attn_sharedkv", + TORCH_FN(xllm::kernel::npu::sparse_attn_sharedkv)); + m.impl("sparse_attn_sharedkv_metadata", + TORCH_FN(xllm::kernel::npu::sparse_attn_sharedkv_metadata)); + m.impl("quant_lightning_indexer", + TORCH_FN(xllm::kernel::npu::quant_lightning_indexer)); + m.impl("quant_lightning_indexer_metadata", + TORCH_FN(xllm::kernel::npu::quant_lightning_indexer_metadata)); } diff --git a/xllm/core/runtime/py_attention_metadata.cpp b/xllm/core/runtime/py_attention_metadata.cpp index 9626a6fe64..1999a6a177 100644 --- a/xllm/core/runtime/py_attention_metadata.cpp +++ b/xllm/core/runtime/py_attention_metadata.cpp @@ -26,6 +26,24 @@ limitations under the License. namespace py = pybind11; namespace xllm { +namespace { + +struct PythonObjectHolder final { + explicit PythonObjectHolder(py::object value) : value(std::move(value)) {} + + ~PythonObjectHolder() { + if (!Py_IsInitialized()) { + (void)value.release(); + return; + } + py::gil_scoped_acquire gil; + value = py::object(); + } + + py::object value; +}; + +} // namespace void register_attention_metadata_views(py::module_& module) { py::class_(module, "ExpandedDecodeMetadataView") @@ -70,6 +88,8 @@ void register_attention_metadata_views(py::module_& module) { &PyAttentionMetadataView::kv_seq_lens_host_values) .def_property_readonly("q_seq_lens_host", &PyAttentionMetadataView::q_seq_lens_host) + .def_property_readonly("multi_block_tables", + &PyAttentionMetadataView::multi_block_tables) .def_property_readonly("block_table", &PyAttentionMetadataView::block_table) .def_property_readonly("kv_seq_lens", @@ -83,6 +103,31 @@ void register_attention_metadata_views(py::module_& module) { .def_property_readonly("q_seq_lens", &PyAttentionMetadataView::q_seq_lens) .def_property_readonly("expanded_decode_metadata", &PyAttentionMetadataView::expanded_decode_metadata) + .def_property_readonly("max_query_len", + &PyAttentionMetadataView::max_query_len) + .def_property_readonly("max_seq_len", + &PyAttentionMetadataView::max_seq_len) + .def_property("dsa_metadata", + &PyAttentionMetadataView::dsa_metadata, + &PyAttentionMetadataView::set_dsa_metadata) + .def_property("dsa_positions", + &PyAttentionMetadataView::dsa_positions, + &PyAttentionMetadataView::set_dsa_positions) + .def_property("dsa_cos_sin", + &PyAttentionMetadataView::dsa_cos_sin, + &PyAttentionMetadataView::set_dsa_cos_sin) + .def_property("dsa_c4_cos_sin", + &PyAttentionMetadataView::dsa_c4_cos_sin, + &PyAttentionMetadataView::set_dsa_c4_cos_sin) + .def_property("dsa_c128_cos_sin", + &PyAttentionMetadataView::dsa_c128_cos_sin, + &PyAttentionMetadataView::set_dsa_c128_cos_sin) + .def_property("dsa_graph_block_table_cols", + &PyAttentionMetadataView::dsa_graph_block_table_cols, + &PyAttentionMetadataView::set_dsa_graph_block_table_cols) + .def_property("dsa_graph_mode", + &PyAttentionMetadataView::dsa_graph_mode, + &PyAttentionMetadataView::set_dsa_graph_mode) .def_property_readonly("is_prefill", &PyAttentionMetadataView::is_prefill) .def_property_readonly("is_chunked_prefill", &PyAttentionMetadataView::is_chunked_prefill); @@ -158,6 +203,7 @@ PyAttentionMetadataView::PyAttentionMetadataView( std::shared_ptr metadata, const ModelInputParams& params) : PyAttentionMetadataView(std::move(metadata)) { + multi_block_tables_ = params.multi_block_tables; linear_state_indices_ = params.embedding.linear_state_indices; dp_token_counts_ = params.parallel.raw_dp_global_token_nums.empty() ? params.parallel.dp_global_token_nums @@ -232,11 +278,93 @@ py::object PyAttentionMetadataView::q_seq_lens_host() const { return optional_tensor(q_seq_lens_host_); } +py::list PyAttentionMetadataView::multi_block_tables() const { + py::list tables; + for (const torch::Tensor& table : multi_block_tables_) { + tables.append(optional_tensor(table)); + } + return tables; +} + PyExpandedDecodeMetadataView PyAttentionMetadataView::expanded_decode_metadata() const { return PyExpandedDecodeMetadataView(metadata_); } +int64_t PyAttentionMetadataView::max_query_len() const { + return metadata_->max_query_len; +} + +int64_t PyAttentionMetadataView::max_seq_len() const { + return metadata_->max_seq_len; +} + +py::object PyAttentionMetadataView::dsa_metadata() const { + if (!dsa_metadata_holder_) { + return py::none(); + } + return std::static_pointer_cast(dsa_metadata_holder_) + ->value; +} + +void PyAttentionMetadataView::set_dsa_metadata(py::object value) { + if (value.is_none()) { + dsa_metadata_holder_.reset(); + return; + } + dsa_metadata_holder_ = std::make_shared(std::move(value)); +} + +py::object PyAttentionMetadataView::dsa_positions() const { + return optional_tensor(dsa_positions_); +} + +void PyAttentionMetadataView::set_dsa_positions(py::object value) { + dsa_positions_ = + value.is_none() ? torch::Tensor() : value.cast(); +} + +py::object PyAttentionMetadataView::dsa_cos_sin() const { + return optional_tensor(dsa_cos_sin_); +} + +void PyAttentionMetadataView::set_dsa_cos_sin(py::object value) { + dsa_cos_sin_ = + value.is_none() ? torch::Tensor() : value.cast(); +} + +py::object PyAttentionMetadataView::dsa_c4_cos_sin() const { + return optional_tensor(dsa_c4_cos_sin_); +} + +void PyAttentionMetadataView::set_dsa_c4_cos_sin(py::object value) { + dsa_c4_cos_sin_ = + value.is_none() ? torch::Tensor() : value.cast(); +} + +py::object PyAttentionMetadataView::dsa_c128_cos_sin() const { + return optional_tensor(dsa_c128_cos_sin_); +} + +void PyAttentionMetadataView::set_dsa_c128_cos_sin(py::object value) { + dsa_c128_cos_sin_ = + value.is_none() ? torch::Tensor() : value.cast(); +} + +int64_t PyAttentionMetadataView::dsa_graph_block_table_cols() const { + return dsa_graph_block_table_cols_; +} + +void PyAttentionMetadataView::set_dsa_graph_block_table_cols(int64_t value) { + dsa_graph_block_table_cols_ = value; +} + +bool PyAttentionMetadataView::dsa_graph_mode() const { return dsa_graph_mode_; } + +void PyAttentionMetadataView::set_dsa_graph_mode(bool value) { + dsa_graph_mode_ = value; +} + bool PyAttentionMetadataView::is_prefill() const { return metadata_->is_prefill; } diff --git a/xllm/core/runtime/py_attention_metadata.h b/xllm/core/runtime/py_attention_metadata.h index e0a86dd70c..4573f77b38 100644 --- a/xllm/core/runtime/py_attention_metadata.h +++ b/xllm/core/runtime/py_attention_metadata.h @@ -71,6 +71,7 @@ class PyAttentionMetadataView final { pybind11::object kv_seq_lens_host() const; const std::vector& kv_seq_lens_host_values() const; pybind11::object q_seq_lens_host() const; + pybind11::list multi_block_tables() const; pybind11::object block_table() const; pybind11::object kv_seq_lens() const; pybind11::object linear_state_indices() const; @@ -78,6 +79,22 @@ class PyAttentionMetadataView final { const std::vector& dp_token_counts() const; pybind11::object q_seq_lens() const; PyExpandedDecodeMetadataView expanded_decode_metadata() const; + int64_t max_query_len() const; + int64_t max_seq_len() const; + pybind11::object dsa_metadata() const; + void set_dsa_metadata(pybind11::object value); + pybind11::object dsa_positions() const; + void set_dsa_positions(pybind11::object value); + pybind11::object dsa_cos_sin() const; + void set_dsa_cos_sin(pybind11::object value); + pybind11::object dsa_c4_cos_sin() const; + void set_dsa_c4_cos_sin(pybind11::object value); + pybind11::object dsa_c128_cos_sin() const; + void set_dsa_c128_cos_sin(pybind11::object value); + int64_t dsa_graph_block_table_cols() const; + void set_dsa_graph_block_table_cols(int64_t value); + bool dsa_graph_mode() const; + void set_dsa_graph_mode(bool value); bool is_prefill() const; bool is_chunked_prefill() const; @@ -90,8 +107,16 @@ class PyAttentionMetadataView final { std::shared_ptr metadata_; torch::Tensor kv_seq_lens_host_; torch::Tensor q_seq_lens_host_; + std::vector multi_block_tables_; torch::Tensor linear_state_indices_; std::vector dp_token_counts_; + std::shared_ptr dsa_metadata_holder_; + torch::Tensor dsa_positions_; + torch::Tensor dsa_cos_sin_; + torch::Tensor dsa_c4_cos_sin_; + torch::Tensor dsa_c128_cos_sin_; + int64_t dsa_graph_block_table_cols_ = 0; + bool dsa_graph_mode_ = false; }; } // namespace xllm diff --git a/xllm/core/runtime/py_executor_impl.cpp b/xllm/core/runtime/py_executor_impl.cpp index 7dac2df748..d1d6fc7e50 100644 --- a/xllm/core/runtime/py_executor_impl.cpp +++ b/xllm/core/runtime/py_executor_impl.cpp @@ -41,10 +41,22 @@ namespace py = pybind11; namespace xllm { namespace { +// Python collective wrappers call back into the C++ ProcessGroups owned by +// the model. The executor is single-threaded per worker, but thread-local +// storage keeps concurrent workers isolated. +thread_local PyCausalLM* active_py_causal_lm = nullptr; + py::object optional_tensor(const torch::Tensor& tensor) { return tensor.defined() ? py::cast(tensor) : py::none(); } +py::object optional_tensor(const std::optional& tensor) { + if (!tensor.has_value() || !tensor->defined()) { + return py::none(); + } + return py::cast(*tensor); +} + void clear_python_object(py::object& object) { if (!object) { return; @@ -62,6 +74,34 @@ void clear_python_object(py::object& object) { PYBIND11_EMBEDDED_MODULE(xllm_runtime, m) { register_attention_metadata_views(m); + // Reuse native C++ process groups instead of creating a second HCCL + // communicator from Python. These functions are valid only during a model + // forward, while PyExecutorImpl has set active_py_causal_lm. + m.def("tp_all_reduce", [](torch::Tensor tensor) { + if (active_py_causal_lm != nullptr) { + active_py_causal_lm->tp_all_reduce(tensor); + } + return tensor; + }); + m.def("tp_all_gather", [](torch::Tensor tensor, int64_t dim) { + if (active_py_causal_lm != nullptr) { + return active_py_causal_lm->tp_all_gather(tensor, dim); + } + return tensor; + }); + m.def("moe_tp_all_reduce", [](torch::Tensor tensor) { + if (active_py_causal_lm != nullptr) { + active_py_causal_lm->moe_tp_all_reduce(tensor); + } + return tensor; + }); + m.def("moe_ep_all_reduce", [](torch::Tensor tensor) { + if (active_py_causal_lm != nullptr) { + active_py_causal_lm->moe_ep_all_reduce(tensor); + } + return tensor; + }); + #if defined(USE_NPU) py::class_>(m, "LayerSynchronizer") @@ -95,7 +135,12 @@ PyExecutorImpl::PyExecutorImpl(CausalLM* model, options_.max_seqs_per_batch()); } -PyExecutorImpl::~PyExecutorImpl() { clear_python_object(py_executor_); } +PyExecutorImpl::~PyExecutorImpl() { + if (active_py_causal_lm == py_causal_lm_) { + active_py_causal_lm = nullptr; + } + clear_python_object(py_executor_); +} ForwardInput PyExecutorImpl::prepare_inputs(Batch& batch) { return batch.prepare_forward_input( @@ -108,6 +153,9 @@ ModelOutput PyExecutorImpl::run(const torch::Tensor& tokens, const ModelInputParams& params) { torch::NoGradGuard no_grad; COUNTER_INC(num_model_execution_total_eager); + // Keep this active through the subsequent PyCausalLM::logits() call: the + // sharded lm_head performs its TP gather after execute() returns. + active_py_causal_lm = py_causal_lm_; // Build or reuse attention metadata. std::shared_ptr attn_metadata = @@ -126,11 +174,21 @@ ModelOutput PyExecutorImpl::run(const torch::Tensor& tokens, py::list kv_caches_py; for (auto& kv : kv_caches) { // Slot order must match ``LayerCache`` on the Python side. - kv_caches_py.append(py::make_tuple(optional_tensor(kv.get_k_cache()), - optional_tensor(kv.get_v_cache()), - optional_tensor(kv.get_index_cache()), - optional_tensor(kv.get_conv_cache()), - optional_tensor(kv.get_ssm_cache()))); + // Keep this order synchronized with LayerCache/_LAYER_CACHE_SLOTS. + // Generic caches use the first five entries; DeepSeek-V4 uses the + // trailing six entries returned by KVCache's DSV4 getters. + kv_caches_py.append( + py::make_tuple(optional_tensor(kv.get_k_cache()), + optional_tensor(kv.get_v_cache()), + optional_tensor(kv.get_index_cache()), + optional_tensor(kv.get_conv_cache()), + optional_tensor(kv.get_ssm_cache()), + optional_tensor(kv.get_swa_cache()), + optional_tensor(kv.get_compress_kv_state()), + optional_tensor(kv.get_compress_score_state()), + optional_tensor(kv.get_compress_index_kv_state()), + optional_tensor(kv.get_compress_index_score_state()), + optional_tensor(kv.get_indexer_cache_scale()))); } py_executor_.attr("bind_kv_caches")(kv_caches_py); kv_bound_ = true; diff --git a/xllm/models/llm/py_causal_lm.cpp b/xllm/models/llm/py_causal_lm.cpp index 370774f9ac..896c5765b4 100644 --- a/xllm/models/llm/py_causal_lm.cpp +++ b/xllm/models/llm/py_causal_lm.cpp @@ -27,6 +27,7 @@ limitations under the License. #include "core/framework/config/execution_config.h" #include "core/framework/model/model_output.h" #include "core/framework/model_loader.h" +#include "core/framework/parallel_state/process_group.h" #include "core/framework/state_dict/state_dict.h" #include "models/py_model_helper.h" @@ -91,86 +92,86 @@ PyCausalLM::PyCausalLM(const ModelContext& context) << "Python models support only ep_size=1 or ep_size=world_size."; CHECK(parallel_args.moe_tp_group_ != nullptr); - ProcessGroup* moe_tp_group = parallel_args.moe_tp_group_; - ProcessGroup* ep_group = nullptr; + moe_tp_group_ = parallel_args.moe_tp_group_; if (ep_size_ > 1) { CHECK(parallel_args.moe_ep_group_ != nullptr); - ep_group = parallel_args.moe_ep_group_; + moe_ep_group_ = parallel_args.moe_ep_group_; } - moe_tp_size_ = (moe_tp_group != nullptr) ? moe_tp_group->world_size() : 1; - moe_tp_rank_ = (moe_tp_group != nullptr) ? moe_tp_group->rank() : 0; - ep_rank_ = (ep_group != nullptr) ? ep_group->rank() : 0; + moe_tp_size_ = (moe_tp_group_ != nullptr) ? moe_tp_group_->world_size() : 1; + moe_tp_rank_ = (moe_tp_group_ != nullptr) ? moe_tp_group_->rank() : 0; + ep_rank_ = (moe_ep_group_ != nullptr) ? moe_ep_group_->rank() : 0; py::gil_scoped_acquire gil; - py::object init_process_group = - py::module_::import("xllm.python.distributed").attr("init_process_group"); - CHECK(!parallel_args.python_rendezvous_host_.empty()); - CHECK_GT(parallel_args.python_rendezvous_port_, 0); - const int32_t global_rank = parallel_args.rank(); - const int32_t global_world_size = parallel_args.world_size(); - if (tp_size_ > 1) { - init_process_group("tp", - parallel_args.python_rendezvous_host_, - parallel_args.python_rendezvous_port_, - tp_rank_, - tp_size_, - c10::str(device_), - global_rank, - global_world_size, - global_rank / tp_size_); - } - if (dp_size_ > 1) { - init_process_group("dp", - parallel_args.python_rendezvous_host_, - parallel_args.python_rendezvous_port_, - dp_rank_, - dp_size_, - c10::str(device_), - global_rank, - global_world_size, - global_rank % tp_size_); - } - if (moe_tp_size_ > 1) { - init_process_group("moe_tp", - parallel_args.python_rendezvous_host_, - parallel_args.python_rendezvous_port_, - moe_tp_rank_, - moe_tp_size_, - c10::str(device_), - global_rank, - global_world_size, - global_rank / moe_tp_size_); - } - if (ep_size_ > 1) { - init_process_group("moe_ep", - parallel_args.python_rendezvous_host_, - parallel_args.python_rendezvous_port_, - ep_rank_, - ep_size_, - c10::str(device_), - global_rank, - global_world_size, - global_rank % moe_tp_size_); - } - if (cp_size_ > 1) { - // CP shards sequence tokens; its group is strided by tp_size -- ranks with - // the same (dp, tp) slot but different cp_rank. The group index selects - // that (dp, tp) slot: dp block (global_rank / (cp_size*tp_size)) times - // tp_size, plus the tp offset within it. TP and CP are orthogonal, so both - // groups may be initialized on the same device off the shared rendezvous - // endpoint. - const int32_t cp_group_index = - (global_rank / (cp_size_ * tp_size_)) * tp_size_ + - global_rank % tp_size_; - init_process_group("cp", - parallel_args.python_rendezvous_host_, - parallel_args.python_rendezvous_port_, - cp_rank_, - cp_size_, - c10::str(device_), - global_rank, - global_world_size, - cp_group_index); + // DeepSeek-V4 uses the native ProcessGroups exposed by xllm_runtime. This + // avoids creating a second HCCL communicator from Python on the same NPU. + // Other Python models retain the existing c10d setup for their DP/CP paths; + // their TP/MoE layers can still use the bridge when available. + if (model_args_.model_type() != "deepseek_v4") { + py::object init_process_group = + py::module_::import("xllm.python.distributed") + .attr("init_process_group"); + CHECK(!parallel_args.python_rendezvous_host_.empty()); + CHECK_GT(parallel_args.python_rendezvous_port_, 0); + const int32_t global_rank = parallel_args.rank(); + const int32_t global_world_size = parallel_args.world_size(); + if (tp_size_ > 1) { + init_process_group("tp", + parallel_args.python_rendezvous_host_, + parallel_args.python_rendezvous_port_, + tp_rank_, + tp_size_, + c10::str(device_), + global_rank, + global_world_size, + global_rank / tp_size_); + } + if (dp_size_ > 1) { + init_process_group("dp", + parallel_args.python_rendezvous_host_, + parallel_args.python_rendezvous_port_, + dp_rank_, + dp_size_, + c10::str(device_), + global_rank, + global_world_size, + global_rank % tp_size_); + } + if (moe_tp_size_ > 1) { + init_process_group("moe_tp", + parallel_args.python_rendezvous_host_, + parallel_args.python_rendezvous_port_, + moe_tp_rank_, + moe_tp_size_, + c10::str(device_), + global_rank, + global_world_size, + global_rank / moe_tp_size_); + } + if (ep_size_ > 1) { + init_process_group("moe_ep", + parallel_args.python_rendezvous_host_, + parallel_args.python_rendezvous_port_, + ep_rank_, + ep_size_, + c10::str(device_), + global_rank, + global_world_size, + global_rank % moe_tp_size_); + } + if (cp_size_ > 1) { + const int32_t cp_group_index = + (global_rank / (cp_size_ * tp_size_)) * tp_size_ + + global_rank % tp_size_; + init_process_group("cp", + parallel_args.python_rendezvous_host_, + parallel_args.python_rendezvous_port_, + cp_rank_, + cp_size_, + c10::str(device_), + global_rank, + global_world_size, + cp_group_index); + } } const std::string module_name = context.get_model_args().model_type().empty() ? std::string("Qwen3ForCausalLM") @@ -249,6 +250,56 @@ torch::Tensor PyCausalLM::logits(const torch::Tensor& hidden_states, return out.cast(); } +void PyCausalLM::tp_all_reduce(torch::Tensor& tensor) { + if (tp_group_ != nullptr) { + tp_group_->allreduce(tensor); + } +} + +torch::Tensor PyCausalLM::tp_all_gather(const torch::Tensor& tensor, + int64_t dim) { + if (tp_group_ == nullptr) { + return tensor; + } + auto gathered = tp_group_->allgather_base_sync(tensor); + const int64_t world_size = tp_group_->world_size(); + const int64_t ndim = tensor.dim(); + if (dim < 0) { + dim += ndim; + } + TORCH_CHECK(dim >= 0 && dim < ndim, + "tensor-parallel gather dimension out of range: ", + dim); + + // allgather_base_sync returns [world_size, *input_shape]. Move the leading + // world dimension next to the requested dimension, then merge the pair. + std::vector permutation; + permutation.reserve(static_cast(ndim + 1)); + for (int64_t index = 1; index <= dim; ++index) { + permutation.push_back(index); + } + permutation.push_back(0); + for (int64_t index = dim + 1; index < ndim + 1; ++index) { + permutation.push_back(index); + } + gathered = gathered.permute(permutation); + auto output_shape = tensor.sizes().vec(); + output_shape[dim] *= world_size; + return gathered.reshape(output_shape).contiguous(); +} + +void PyCausalLM::moe_tp_all_reduce(torch::Tensor& tensor) { + if (moe_tp_group_ != nullptr) { + moe_tp_group_->allreduce(tensor); + } +} + +void PyCausalLM::moe_ep_all_reduce(torch::Tensor& tensor) { + if (moe_ep_group_ != nullptr) { + moe_ep_group_->allreduce(tensor); + } +} + 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 ae9913b6d5..20d9203ef1 100644 --- a/xllm/models/llm/py_causal_lm.h +++ b/xllm/models/llm/py_causal_lm.h @@ -78,6 +78,14 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { bool share_weights_from(CausalLM& source) override; + // Reuse the native process groups from the Python model path. Python falls + // back to its c10d implementation only when this embedded bridge is absent + // (for example, in pure-Python unit tests). + void tp_all_reduce(torch::Tensor& tensor); + torch::Tensor tp_all_gather(const torch::Tensor& tensor, int64_t dim); + void moe_tp_all_reduce(torch::Tensor& tensor); + void moe_ep_all_reduce(torch::Tensor& tensor); + pybind11::object& python_model() { return py_model_; } const pybind11::object& config_dict() const { return config_dict_; } @@ -100,6 +108,8 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { int64_t cp_size_ = 1; int64_t cp_rank_ = 0; ProcessGroup* tp_group_ = nullptr; + ProcessGroup* moe_tp_group_ = nullptr; + ProcessGroup* moe_ep_group_ = nullptr; pybind11::object py_model_; pybind11::object config_dict_; diff --git a/xllm/python/attention/backend.py b/xllm/python/attention/backend.py index bf6ab94c47..349b1cadf9 100644 --- a/xllm/python/attention/backend.py +++ b/xllm/python/attention/backend.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: from xllm.python.layers.attention import Attention + @dataclass(frozen=True, slots=True) class LayerCache: """Every cache a layer may own, named rather than positional. @@ -42,10 +43,30 @@ class LayerCache: index: torch.Tensor | None = None conv: torch.Tensor | None = None ssm: torch.Tensor | None = None + # DeepSeek-V4 DSA cache slots. Generic models leave these as None; the + # tuple order is shared with PyExecutorImpl::bind_kv_caches. + swa: torch.Tensor | None = None + compress_kv_state: torch.Tensor | None = None + compress_score_state: torch.Tensor | None = None + compress_index_kv_state: torch.Tensor | None = None + compress_index_score_state: torch.Tensor | None = None + indexer_scale: torch.Tensor | None = None #: Field order of the tuple form, which is what the C++ executor hands over. -_LAYER_CACHE_SLOTS = ("key", "value", "index", "conv", "ssm") +_LAYER_CACHE_SLOTS = ( + "key", + "value", + "index", + "conv", + "ssm", + "swa", + "compress_kv_state", + "compress_score_state", + "compress_index_kv_state", + "compress_index_score_state", + "indexer_scale", +) LayerCacheInput = LayerCache | tuple[torch.Tensor | None, ...] @@ -58,12 +79,8 @@ def normalize_layer_caches(caches: Sequence[LayerCacheInput]) -> list[LayerCache normalized.append(cache) continue if not 2 <= len(cache) <= len(_LAYER_CACHE_SLOTS): - raise ValueError( - "layer cache must hold between K/V and " - f"{'/'.join(_LAYER_CACHE_SLOTS)} tensors" - ) - slots = [None if tensor is None or not tensor.numel() else tensor - for tensor in cache] + raise ValueError(f"layer cache must hold between K/V and {'/'.join(_LAYER_CACHE_SLOTS)} tensors") + slots = [None if tensor is None or not tensor.numel() else tensor for tensor in cache] slots.extend([None] * (len(_LAYER_CACHE_SLOTS) - len(slots))) normalized.append(LayerCache(*slots)) return normalized @@ -84,6 +101,16 @@ class AttentionMetadata(Protocol): paged_kv_last_page_len_host: torch.Tensor | None block_table: torch.Tensor | None kv_seq_lens: torch.Tensor | None + max_query_len: int + max_seq_len: int + multi_block_tables: Sequence[torch.Tensor | None] + dsa_metadata: object | None + dsa_positions: torch.Tensor | None + dsa_cos_sin: torch.Tensor | None + dsa_c4_cos_sin: torch.Tensor | None + dsa_c128_cos_sin: torch.Tensor | None + dsa_graph_block_table_cols: int + dsa_graph_mode: bool linear_state_indices: torch.Tensor | None has_initial_state: torch.Tensor | None dp_token_counts: Sequence[int] @@ -111,6 +138,25 @@ class MlaIndexContext: update_index_cache: Callable[[torch.Tensor], None] +@dataclass(frozen=True) +class DsaIndexContext: + """Per-forward cache and metadata view consumed by the DSV4 indexer.""" + + index_cache: torch.Tensor + indexer_scale: torch.Tensor | None + slot_mapping: torch.Tensor + block_table: torch.Tensor | None + cmp_block_table: torch.Tensor | None + kv_state: torch.Tensor | None + score_state: torch.Tensor | None + kv_block_table: torch.Tensor | None + score_block_table: torch.Tensor | None + actual_seq_q: torch.Tensor + actual_seq_kv: torch.Tensor + start_pos: torch.Tensor | None + qli_metadata: torch.Tensor | None + + class AttentionBackend(ABC): @abstractmethod def bind_kv_caches(self, kv_caches: list[LayerCache]) -> None: @@ -131,7 +177,7 @@ def execute( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - layer: "Attention", + layer: Attention, ) -> torch.Tensor: pass @@ -151,7 +197,7 @@ def execute_mla( q_pe: torch.Tensor, k_latent: torch.Tensor, k_pe: torch.Tensor, - layer: "Attention", + layer: Attention, topk: torch.Tensor | None = None, ) -> torch.Tensor: """Absorbed-MLA attention over paged latent (nope) + rope caches. @@ -161,11 +207,9 @@ def execute_mla( LightningIndexer; otherwise a dense MLA path is requested. Backends that do not implement MLA raise. """ - raise NotImplementedError( - f"{type(self).__name__} does not support MLA" - ) + raise NotImplementedError(f"{type(self).__name__} does not support MLA") - def mla_index_context(self, layer: "Attention") -> MlaIndexContext: + def mla_index_context(self, layer: Attention) -> MlaIndexContext: """Public hook for an optional LightningIndexer. Hands out the paged index cache view (``LayerCache.index``) plus the @@ -173,6 +217,4 @@ def mla_index_context(self, layer: "Attention") -> MlaIndexContext: touches ``backend._metadata`` / ``backend._kv_caches`` directly. Backends that do not support the sparse MLA indexer raise. """ - raise NotImplementedError( - f"{type(self).__name__} does not support MLA indexer" - ) + raise NotImplementedError(f"{type(self).__name__} does not support MLA indexer") diff --git a/xllm/python/attention/dsa_attention.py b/xllm/python/attention/dsa_attention.py new file mode 100644 index 0000000000..2b9c41a518 --- /dev/null +++ b/xllm/python/attention/dsa_attention.py @@ -0,0 +1,778 @@ +# 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. + +"""DeepSeek-V4 DSA attention backend. + +Consumes the :class:`DsaMetadata` built by :mod:`dsa_metadata` and drives the +two-stage sparse attention (``sparse_attn_sharedkv``), the KV compressor, and +the quantized lightning indexer. This is the Python-path counterpart of the +C++ ``DSAttentionImpl`` (core/layers/npu_torch/deepseek_sparse_attention.cpp). + +The backend owns no KV storage: caches are bound from the C++ executor's +``LayerCache`` 11-tuple. Per step it (1) builds DSA metadata from the framework +``multi_block_tables``, (2) resolves the per-layer 8-cache mapping, (3) writes +new KV into the SWA cache, (4) runs the compressor into the compressed cache +when ``compress_ratio > 1``, (5) runs the indexer to pick top-k compressed +blocks when ``compress_ratio == 4``, and (6) calls ``sparse_attn_sharedkv``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from xllm.python.attention.backend import ( + AttentionBackend, + DsaIndexContext, + LayerCache, +) +from xllm.python.attention.dsa_metadata import ( + DSA_CACHE_SLIDING_WINDOW, + DSA_CACHE_TOKEN, + DsaMetadata, + DsaMetadataBuilder, + build_cache_specs, +) +from xllm.python.model_executor.forward_context import get_forward_context + +if TYPE_CHECKING: + from xllm.python.attention.backend import AttentionMetadata + from xllm.python.layers.attention import Attention + +# Sparse mask modes used by the C++ DSA attention (rightDownCausal variants). +_MASK_MODE_RIGHT_DOWN_CAUSAL = 3 +_MASK_MODE_COMPRESS = 4 + + +@dataclass +class _DsaCacheMapping: + """Per-layer resolved cache indices (mirrors C++ ``DsaCacheMapping``).""" + + cmp_cache_idx: int = -1 + index_cache_idx: int = -1 + indexer_scale_cache_idx: int = -1 + ori_cache_idx: int = -1 + kv_state_cache_idx: int = -1 + score_state_cache_idx: int = -1 + index_kv_state_cache_idx: int = -1 + index_score_state_cache_idx: int = -1 + + +@dataclass(frozen=True) +class _DsaForwardMeta: + """Subset of C++ ModelInputParams::meta used by DSV4 metadata builders.""" + + q_max_seq_len: int + kv_max_seq_len: int + + +class DsaAttentionBackend(AttentionBackend): + """DSA attention backend for DeepSeek-V4 on NPU. + + The model supplies its config so the backend can rebuild the static + ``caches_info`` / ``group_infos`` once (mirroring + ``deepseek_v4_build_cache_specs``) and precompute the per-step RoPE / Hadamard + tables the indexer needs. + """ + + def __init__( + self, + compress_ratios: list[int], + window_size: int, + n_layers: int, + num_heads: int, + attn_head_dim: int, + index_topk: int, + index_n_heads: int, + index_head_dim: int, + rope_head_dim: int, + device: torch.device, + dtype: torch.dtype, + ) -> None: + self.caches_info, self.group_infos = build_cache_specs(compress_ratios, window_size, n_layers) + self._builder = DsaMetadataBuilder(self.caches_info, self.group_infos) + self.window_size = window_size + self.index_topk = index_topk + self.index_n_heads = index_n_heads + self.index_head_dim = index_head_dim + self.rope_head_dim = rope_head_dim + self.num_heads = num_heads + self.head_dim = attn_head_dim + self.device = device + self.dtype = dtype + self.scale = attn_head_dim**-0.5 + + self._kv_caches: list[LayerCache] = [] + self._metadata: AttentionMetadata | None = None + + # -- AttentionBackend interface ----------------------------------------- + + def bind_kv_caches(self, kv_caches: list[LayerCache]) -> None: + self._kv_caches = kv_caches + + def _current_forward_metadata(self) -> AttentionMetadata: + try: + return get_forward_context().metadata + except RuntimeError: + if self._metadata is None: + raise + return self._metadata + + def prepare( + self, + metadata: AttentionMetadata, + *, + graph_mode: bool = False, + ) -> None: + if graph_mode: + raise NotImplementedError("DeepSeek-V4 ACL graph support is not part of the eager DSA backend") + self._metadata = metadata + + def prepare_dsa_metadata_for_forward( + self, + metadata: AttentionMetadata | None = None, + ) -> None: + """Build DSA metadata inside model forward, matching C++ ownership/order.""" + metadata = metadata or self._metadata + assert metadata is not None + multi_block_tables = list(metadata.multi_block_tables) + kv_seq_lens_host = metadata.kv_seq_lens_host + kv_seq_lens = ( + kv_seq_lens_host.cpu().tolist() if kv_seq_lens_host is not None and kv_seq_lens_host.numel() > 0 else [] + ) + q_seq_lens_host = getattr(metadata, "q_seq_lens_host", None) + q_seq_lens = ( + q_seq_lens_host.cpu().tolist() if q_seq_lens_host is not None and q_seq_lens_host.numel() > 0 else None + ) + # DSA RoPE tables and positions are model-owned; the backend reads them + # off the metadata when the model attaches them (see attach_rope_tables). + positions = getattr(metadata, "dsa_positions", None) + if positions is None: + positions = torch.empty(0, dtype=torch.int64) + dsa_cos_sin = getattr(metadata, "dsa_cos_sin", None) + dsa_metadata = self._builder.build( + multi_block_tables=multi_block_tables, + kv_seq_lens=kv_seq_lens, + q_seq_lens=q_seq_lens, + positions=positions, + dsa_cos_sin=dsa_cos_sin, + is_prefill=metadata.is_prefill, + is_chunked_prefill=metadata.is_chunked_prefill, + enable_graph=False, + ) + self._populate_dsa_rope(dsa_metadata, metadata) + self._move_metadata_to_device(dsa_metadata) + self._build_precomputed_metadata(dsa_metadata, metadata) + metadata.dsa_metadata = dsa_metadata + + def select_dsa_layer_rope( + self, + layer_id: int, + cos_sin_cache: torch.Tensor, + metadata: AttentionMetadata | None = None, + ) -> None: + """Select the main q/kv RoPE group for the current DSV4 layer. + + C++ updates ``DSAMetadata::layer_id/cos/sin`` in the model layer loop + from ``input_rope_by_ratio``. Python keeps the full cache here because + the model and indexer gather it with the current input positions, but + the selected group and lifetime are otherwise identical. + """ + metadata = metadata or self._metadata + if metadata is None or metadata.dsa_metadata is None: + raise RuntimeError("DSA metadata must be prepared before selecting layer RoPE") + dsa = metadata.dsa_metadata + chunks = cos_sin_cache.chunk(2, dim=-1) + dsa.layer_id = layer_id + dsa.cos_table = chunks[0].contiguous() + dsa.sin_table = chunks[1].contiguous() + + def _populate_dsa_rope( + self, + dsa: DsaMetadata, + metadata: AttentionMetadata | None = None, + ) -> None: + """Build request-shaped RoPE tensors for the current forward.""" + metadata = metadata or self._metadata + css = getattr(metadata, "dsa_cos_sin", None) if metadata is not None else None + if dsa.cos_table is None and css is not None and css.numel() > 0: + dsa.cos_table, dsa.sin_table = (tensor.contiguous() for tensor in css.chunk(2, dim=-1)) + c4css = getattr(metadata, "dsa_c4_cos_sin", None) if metadata is not None else None + if c4css is not None and dsa.c4_pad_positions.numel() > 0: + c4_idx = dsa.c4_pad_positions.clamp_min(0).long().to(c4css.device) + dsa.c4_cos, dsa.c4_sin = (tensor.contiguous() for tensor in c4css.index_select(0, c4_idx).chunk(2, dim=-1)) + c128css = getattr(metadata, "dsa_c128_cos_sin", None) if metadata is not None else None + if c128css is not None and dsa.c128_pad_positions.numel() > 0: + c128_idx = dsa.c128_pad_positions.clamp_min(0).long().to(c128css.device) + dsa.c128_cos, dsa.c128_sin = ( + tensor.contiguous() for tensor in c128css.index_select(0, c128_idx).chunk(2, dim=-1) + ) + + def execute( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: Attention, + ) -> torch.Tensor: + """Full DSA attention path for one layer. + + ``q``/``k``/``v`` here are the model-projected, RoPE-applied tensors the + DeepseekV4 attention layer hands in; the backend only owns cache writes + and the kernel dispatch. + """ + metadata = self._current_forward_metadata() + dsa = getattr(metadata, "dsa_metadata", None) + assert dsa is not None + # Late-populate dsa.cos_table/sin_table if prepare ran before the model + # attached the RoPE tables (prepare is called by the executor before + # model.forward, so _dsa_cos_sin may have been None at prepare time). + if dsa.cos_table is None or dsa.sin_table is None: + self._populate_dsa_rope(dsa, metadata) + # Late-populate per-ratio compressed RoPE tables: index the c4/c128 + # compress RoPE cache with the per-token compressed positions + # (c4_pad_positions / c128_pad_positions, built by DsaMetadataBuilder). + # Mirrors C++ DeepseekV4RotaryEmbedding::build(positions_map) per group. + # Also late-populate dsa.input_positions (prepare ran before model.forward + # set self._positions, so it was empty at prepare time). + if dsa.input_positions.numel() == 0: + pos = getattr(metadata, "dsa_positions", None) + if pos is not None and pos.numel() > 0: + dsa.input_positions = pos + if dsa.c4_cos is None or dsa.c128_cos is None: + self._populate_dsa_rope(dsa, metadata) + layer_id = layer.layer_id + compress_ratio = self._layer_compress_ratio(layer_id) + mapping = self._resolve_cache_mapping(layer_id, compress_ratio) + layer_cache = self._kv_caches[layer_id] + is_prefill = metadata.is_prefill + is_chunked_prefill = metadata.is_chunked_prefill + use_temporary_prefill_kv = is_prefill and not is_chunked_prefill + # 1) Prepare ori_kv for attention (mirrors C++ :790-816). + # Full prefill: use kv directly as temporary PA_ND (don't scatter to paged). + # Decode/chunked: scatter to paged SWA cache. + ori_kv = layer_cache.swa + ori_slot = _get_layer_cache_tensor(dsa.slot_mappings, layer_id, mapping.ori_cache_idx) + ori_block_table = _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.ori_cache_idx) + if use_temporary_prefill_kv: + # Prefill: build temporary PA_ND cache from kv (mirrors C++ + # build_prefill_pa_nd_kv, deepseek_sparse_attention.cpp:272-368). + ori_kv_for_attn, ori_block_table_for_attn = _build_prefill_pa_nd_kv( + k, + dsa.actual_seq_lengths_query, + ori_block_table, + self.window_size, + ) + else: + if ori_kv is not None and ori_slot is not None: + _scatter_by_slot(ori_kv, ori_slot, k) + ori_kv_for_attn = ori_kv + ori_block_table_for_attn = ori_block_table + + # 2) Compressor: pool KV into the compressed cache when ratio > 1. + cmp_kv = layer_cache.key + cmp_slot = _get_layer_cache_tensor(dsa.slot_mappings, layer_id, mapping.cmp_cache_idx) + cmp_block_table = _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.cmp_cache_idx) + if compress_ratio > 1 and cmp_kv is not None and cmp_slot is not None: + compressor_fn = getattr(self, "_compressor_fn", None) + if compressor_fn is None: + raise RuntimeError(f"DSA compressor is required for compression ratio {compress_ratio}") + compressed = compressor_fn( + layer_id, + layer_cache, + dsa, + mapping, + cmp_block_table, + compress_ratio, + ) + _scatter_by_slot(cmp_kv, cmp_slot, compressed) + + # 3) Indexer: select top-k compressed blocks when ratio == 4. + compress_topk_idxs: torch.Tensor | None = None + if compress_ratio == 4 and cmp_kv is not None: + indexer_fn = getattr(self, "_indexer_fn", None) + if indexer_fn is None: + raise RuntimeError("DSA indexer is required for compression ratio 4") + compress_topk_idxs = indexer_fn( + layer_id, + layer_cache, + dsa, + mapping, + q, + ) + if compress_topk_idxs is None: + raise RuntimeError("DSA indexer returned no top-k indices") + + # 4) Two-stage sparse attention over original + compressed KV. + # The metadata tensors live on CPU (DsaMetadataBuilder); move to device + # for the NPU kernel, matching the C++ H2D transfer of packed metadata. + if compress_ratio == 1: + sparse_meta = dsa.c1_metadata + elif compress_ratio == 4: + sparse_meta = dsa.c4_metadata + elif compress_ratio == 128: + sparse_meta = dsa.c128_metadata + else: + sparse_meta = None + if sparse_meta is None: + raise RuntimeError(f"DSA sparse metadata is missing for compression ratio {compress_ratio}") + seq_q = dsa.actual_seq_lengths_query + seq_kv = dsa.actual_seq_lengths_kv + sparse_meta_for_kernel = sparse_meta + ori_block_table_for_kernel = ori_block_table_for_attn + cmp_block_table_for_kernel = cmp_block_table + # Match C++ DSAttention's optional contract exactly: prefill and + # chunked prefill pass query cu-seqlens, while decode leaves + # cu_seqlens_ori_kv as std::nullopt. A defined empty tensor selects a + # different ACL optional-input path and causes small decode drift. + use_prefill_attn = is_prefill or is_chunked_prefill + cu_seqlens_ori_kv_for_attn = seq_q if use_prefill_attn else None + out, _lse = _sparse_attn_sharedkv( + q=q, + ori_kv=ori_kv_for_attn, + cmp_kv=cmp_kv if compress_ratio > 1 else None, + ori_sparse_indices=None, + cmp_sparse_indices=compress_topk_idxs, + ori_block_table=ori_block_table_for_kernel, + cmp_block_table=cmp_block_table_for_kernel if compress_ratio > 1 else None, + cu_seqlens_q=seq_q, + cu_seqlens_ori_kv=cu_seqlens_ori_kv_for_attn, + # C++ passes nullopt for compressed KV cu-seqlens; cmp_kv is PA_ND + # and addressed through cmp_block_table/topk. + cu_seqlens_cmp_kv=None, + seqused_q=None, + seqused_kv=seq_kv, + # sinks: the attention sink parameter (attn_sink) is required by the + # sparse_attn_sharedkv kernel (C++ :949 passes attn_sink_ when loaded). + sinks=layer.attn_sink if hasattr(layer, "attn_sink") else None, + metadata=sparse_meta_for_kernel, + softmax_scale=self.scale, + cmp_ratio=compress_ratio, + ori_mask_mode=_MASK_MODE_COMPRESS, + cmp_mask_mode=_MASK_MODE_RIGHT_DOWN_CAUSAL, + ori_win_left=self.window_size - 1, + ori_win_right=0, + layout_q="TND", + layout_kv="PA_ND", + return_softmax_lse=False, + ) + # Full prefill reads a temporary PA_ND cache so attention does not + # depend on the persistent SWA cache. Match C++ step 8 by writing the + # projected KV into the persistent cache only after that attention + # finishes; decode reads this cache on the next forward. + if use_temporary_prefill_kv and ori_kv is not None and ori_slot is not None: + _scatter_by_slot(ori_kv, ori_slot, k) + return out + + def mla_index_context(self, layer: Attention) -> DsaIndexContext: + """Hand the DSA indexer its paged index cache + block tables + slots.""" + metadata = self._current_forward_metadata() + dsa = getattr(metadata, "dsa_metadata", None) + assert dsa is not None + layer_id = layer.layer_id + compress_ratio = self._layer_compress_ratio(layer_id) + mapping = self._resolve_cache_mapping(layer_id, compress_ratio) + layer_cache = self._kv_caches[layer_id] + index_slot = _get_layer_cache_tensor(dsa.slot_mappings, layer_id, mapping.index_cache_idx) + index_block_table = _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.index_cache_idx) + cmp_block_table = _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.cmp_cache_idx) + return DsaIndexContext( + index_cache=layer_cache.index if layer_cache.index is not None else torch.empty(0), + indexer_scale=layer_cache.indexer_scale, + slot_mapping=index_slot if index_slot is not None else torch.empty(0), + block_table=index_block_table, + cmp_block_table=cmp_block_table, + kv_state=layer_cache.compress_kv_state, + score_state=layer_cache.compress_score_state, + kv_block_table=_get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.kv_state_cache_idx), + score_block_table=_get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.score_state_cache_idx), + actual_seq_q=dsa.actual_seq_lengths_query, + actual_seq_kv=dsa.actual_seq_lengths_kv, + start_pos=dsa.start_pos, + qli_metadata=dsa.qli_metadata, + ) + + @property + def num_kv_blocks(self) -> int: + if self._kv_caches and self._kv_caches[0].swa is not None: + return self._kv_caches[0].swa.size(0) + return 0 + + @property + def page_size(self) -> int: + if self._kv_caches and self._kv_caches[0].swa is not None and self._kv_caches[0].swa.dim() > 1: + return self._kv_caches[0].swa.size(1) + return self.window_size + + # -- model-attached state ---------------------------------------------- + # The DeepseekV4 model owns the RoPE tables, Hadamard matrix, and the + # compressor/indexer callables. It attaches them to the backend before the + # first forward so the backend can stage them into DsaMetadata. + + def attach_rope_tables( + self, + positions: torch.Tensor, + dsa_cos_sin: torch.Tensor | None, + graph_bt_cols: int = 0, + c4_cos_sin: torch.Tensor | None = None, + c128_cos_sin: torch.Tensor | None = None, + metadata: AttentionMetadata | None = None, + ) -> None: + metadata = metadata or self._metadata + if metadata is not None: + metadata.dsa_positions = positions + metadata.dsa_cos_sin = dsa_cos_sin + metadata.dsa_c4_cos_sin = c4_cos_sin + metadata.dsa_c128_cos_sin = c128_cos_sin + + def attach_compressor(self, fn) -> None: + """``fn(layer_id, layer_cache, dsa, mapping, cmp_block_table) -> Tensor``.""" + self._compressor_fn = fn + + def attach_indexer(self, fn) -> None: + """``fn(layer_id, layer_cache, dsa, mapping, q) -> Tensor`` (topk idxs).""" + self._indexer_fn = fn + + # -- internals ---------------------------------------------------------- + + def _layer_compress_ratio(self, layer_id: int) -> int: + if layer_id < len(self.caches_info): + caches = self.caches_info[layer_id] + for ci in caches: + if ci.cache_type == DSA_CACHE_TOKEN: + return ci.ratio + return 1 + + def _resolve_cache_mapping(self, layer_id: int, compress_ratio: int) -> _DsaCacheMapping: + """Python port of ``resolve_cache_mapping`` (deepseek_sparse_attention.cpp:92).""" + mapping = _DsaCacheMapping() + if layer_id < 0 or layer_id >= len(self.caches_info): + return mapping + token_ratio_indices: list[int] = [] + swa_indices: list[int] = [] + for cache_idx, ci in enumerate(self.caches_info[layer_id]): + if ci.cache_type == DSA_CACHE_TOKEN and ci.ratio == compress_ratio: + token_ratio_indices.append(cache_idx) + if ci.cache_type == DSA_CACHE_SLIDING_WINDOW: + swa_indices.append(cache_idx) + if token_ratio_indices and compress_ratio > 1: + mapping.cmp_cache_idx = token_ratio_indices[0] + if len(token_ratio_indices) > 1: + mapping.index_cache_idx = token_ratio_indices[1] + if len(token_ratio_indices) > 2: + mapping.indexer_scale_cache_idx = token_ratio_indices[2] + if swa_indices: + mapping.ori_cache_idx = swa_indices[0] + if len(swa_indices) > 1: + mapping.kv_state_cache_idx = swa_indices[1] + if len(swa_indices) > 2: + mapping.score_state_cache_idx = swa_indices[2] + if len(swa_indices) > 3: + mapping.index_kv_state_cache_idx = swa_indices[3] + if len(swa_indices) > 4: + mapping.index_score_state_cache_idx = swa_indices[4] + return mapping + + def _build_precomputed_metadata( + self, + dsa: DsaMetadata, + metadata: AttentionMetadata, + ) -> None: + """Build the AICPU tiling metadata for each compress ratio present. + + Mirrors the C++ ``build_precomputed_metadata`` step: one + ``sparse_attn_sharedkv_metadata`` per ratio, plus one + ``quant_lightning_indexer_metadata`` for the qli path. + """ + from xllm.python import kernels + + seq_q = dsa.actual_seq_lengths_query + seq_kv = dsa.actual_seq_lengths_kv + batch_size = int(max(dsa.actual_seq_lengths_kv.numel(), 1)) + forward_meta = _build_dsa_forward_meta(dsa, metadata) + max_q = forward_meta.q_max_seq_len + max_kv = forward_meta.kv_max_seq_len + is_prefill = max_q > 1 + empty_int32 = torch.empty(0, dtype=torch.int32, device=self.device) + cu_seqlens_ori_kv = seq_q if is_prefill else empty_int32 + cu_seqlens_cmp_kv = empty_int32 + seqused_q = empty_int32 + seqused_kv = seq_kv + # Metadata kernels enqueue asynchronously. Retain their tensor inputs + # on the current forward's DsaMetadata, as C++ DSAMetadata does. + dsa.precomputed_metadata_inputs = tuple( + (seq_q, seq_kv, cu_seqlens_ori_kv, cu_seqlens_cmp_kv, seqused_q, seqused_kv) + ) + for ratio in (1, 4, 128): + has_cmp = ratio > 1 + cmp_topk = self.index_topk if ratio == 4 else 0 + sparse_metadata = kernels.sparse_attn_sharedkv_metadata( + num_heads_q=self.num_heads, + num_heads_kv=1, + head_dim=self.head_dim, + cu_seqlens_q=seq_q, + cu_seqlens_ori_kv=cu_seqlens_ori_kv, + cu_seqlens_cmp_kv=cu_seqlens_cmp_kv, + seqused_q=seqused_q, + seqused_kv=seqused_kv, + batch_size=batch_size, + max_seqlen_q=max_q, + max_seqlen_kv=max_kv, + ori_topk=0, + cmp_topk=cmp_topk, + cmp_ratio=ratio, + ori_mask_mode=_MASK_MODE_COMPRESS, + cmp_mask_mode=_MASK_MODE_RIGHT_DOWN_CAUSAL, + ori_win_left=max(self.window_size - 1, 0), + ori_win_right=0, + layout_q="TND", + layout_kv="PA_ND", + has_ori_kv=True, + has_cmp_kv=has_cmp, + ) + if ratio == 1: + dsa.c1_metadata = sparse_metadata + elif ratio == 4: + dsa.c4_metadata = sparse_metadata + elif ratio == 128: + dsa.c128_metadata = sparse_metadata + query_lens = seq_q[1:].clone() if seq_q.numel() > 1 else dsa.seq_lens_q + key_lens = dsa.seq_lens if dsa.seq_lens.numel() else seq_kv + dsa.precomputed_metadata_inputs += (query_lens, key_lens) + dsa.qli_metadata = kernels.quant_lightning_indexer_metadata( + num_heads_q=max(self.index_n_heads, 1), + num_heads_k=1, + head_dim=max(self.index_head_dim, 1), + query_quant_mode=0, + key_quant_mode=0, + actual_seq_lengths_query=query_lens, + actual_seq_lengths_key=key_lens, + batch_size=int(max(key_lens.size(0), 1)), + max_seqlen_q=max(max_q, 1), + max_seqlen_k=max(max_kv, 1), + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=_MASK_MODE_RIGHT_DOWN_CAUSAL, + pre_tokens=2**63 - 1, + next_tokens=2**63 - 1, + cmp_ratio=4, + device=str(self.device), + ) + + def _move_metadata_to_device(self, dsa: DsaMetadata) -> None: + """Mirror ``deepseek_v4_move_dsa_metadata_to_device`` for eager mode.""" + tensor_fields = ( + "seq_lens", + "seq_lens_q", + "actual_seq_lengths_query", + "actual_seq_lengths_kv", + "kv_cu_seq_lens", + "max_seqlen_q", + "max_seqlen_kv", + "input_positions", + "c4_pad_positions", + "c128_pad_positions", + "start_pos", + "hadamard", + ) + for name in tensor_fields: + tensor = getattr(dsa, name, None) + if tensor is not None: + setattr(dsa, name, tensor.to(self.device)) + for layer_tensors in dsa.block_tables: + for index, tensor in enumerate(layer_tensors): + layer_tensors[index] = tensor.to(self.device) + for layer_tensors in dsa.slot_mappings: + for index, tensor in enumerate(layer_tensors): + layer_tensors[index] = tensor.to(self.device) + + +# --------------------------------------------------------------------------- +# Helpers (faithful ports of C++ free functions). +# --------------------------------------------------------------------------- + + +def _tensor_max_or_zero(tensor: torch.Tensor | None) -> int: + if tensor is None or tensor.numel() == 0: + return 0 + return int(tensor.max().item()) + + +def _build_dsa_forward_meta(dsa: DsaMetadata, metadata: AttentionMetadata) -> _DsaForwardMeta: + """Mirror the C++ max-seqlen inputs used by build_precomputed_metadata. + + C++ computes sparse metadata max sizes from ModelInputParams::meta plus the + host q/kv length vectors: + max(params.meta.q_max_seq_len, max(host.q_seq_lens)) + max(params.meta.kv_max_seq_len, max(host.kv_seq_lens)) + """ + + q_max = int(getattr(metadata, "max_query_len", dsa.max_query_len)) + kv_max = int(getattr(metadata, "max_seq_len", dsa.max_seq_len)) + q_max = max(q_max, _tensor_max_or_zero(getattr(metadata, "q_seq_lens_host", None))) + kv_max = max(kv_max, _tensor_max_or_zero(getattr(metadata, "kv_seq_lens_host", None))) + q_max = max(q_max, int(dsa.max_query_len)) + kv_max = max(kv_max, int(dsa.max_seq_len)) + return _DsaForwardMeta(q_max_seq_len=q_max, kv_max_seq_len=kv_max) + + +def _build_prefill_pa_nd_kv( + kv: torch.Tensor, + cu_seqlens: torch.Tensor, + block_table_hint: torch.Tensor | None, + block_size: int, + cu_seqlens_dst: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Python port of C++ build_prefill_pa_nd_kv (deepseek_sparse_attention.cpp:272-368). + + Builds a temporary PA_ND format KV cache from the current forward's kv + tensor, for full prefill attention (no paged cache needed). + """ + if kv is None or cu_seqlens is None or cu_seqlens.numel() <= 1 or block_size <= 0: + return torch.empty(0), torch.empty(0) + + batch_size = cu_seqlens.numel() - 1 + cu_cpu = cu_seqlens.to(torch.device("cpu")).to(torch.int64) + cu = cu_cpu.tolist() + + dst_cu = None + if cu_seqlens_dst is not None and cu_seqlens_dst.numel() == batch_size + 1: + dst_cu = cu_seqlens_dst.to(torch.device("cpu")).to(torch.int64).tolist() + + # Compute per-request lengths and block counts. + dst_lens = [] + total_blocks = 0 + max_blocks_per_req = 0 + for i in range(batch_size): + q_len = dst_cu[i + 1] - dst_cu[i] if dst_cu is not None else cu[i + 1] - cu[i] + dst_lens.append(q_len) + blocks = (q_len + block_size - 1) // block_size + total_blocks += blocks + max_blocks_per_req = max(max_blocks_per_req, blocks) + + if total_blocks <= 0: + return torch.empty(0), torch.empty(0) + + table_cols = max( + block_table_hint.size(1) if block_table_hint is not None and block_table_hint.dim() > 1 else 0, + max_blocks_per_req, + ) + + # Block 0 is zero-filled padding block; real blocks are 1-based. + packed_kv = torch.zeros( + total_blocks + 1, + block_size, + kv.size(1), + kv.size(2), + dtype=kv.dtype, + device=kv.device, + ) + + table_data = [0] * (batch_size * table_cols) + next_block = 1 + for req in range(batch_size): + q_start = cu[req] + src_len = cu[req + 1] - q_start + q_len = dst_lens[req] + blocks = (q_len + block_size - 1) // block_size + if q_len <= 0 or blocks <= 0: + continue + for j in range(blocks): + table_data[req * table_cols + j] = next_block + j + copy_len = min(q_len, src_len) + if copy_len > 0: + target = packed_kv[next_block : next_block + blocks].view(blocks * block_size, kv.size(1), kv.size(2)) + target[q_len - copy_len : q_len].copy_(kv[q_start : q_start + copy_len]) + next_block += blocks + + table = torch.tensor(table_data, dtype=torch.int32, device=kv.device).view(batch_size, table_cols) + return packed_kv, table + + +def _get_layer_cache_tensor( + layer_tensors: list[list[torch.Tensor]], + layer_id: int, + cache_idx: int, +) -> torch.Tensor | None: + """Python port of ``get_layer_cache_tensor`` (deepseek_sparse_attention.cpp:80).""" + if layer_id < 0 or layer_id >= len(layer_tensors) or cache_idx < 0 or cache_idx >= len(layer_tensors[layer_id]): + return None + return layer_tensors[layer_id][cache_idx] + + +def _scatter_by_slot( + cache: torch.Tensor, + slot_mapping: torch.Tensor, + value: torch.Tensor, +) -> None: + """Python port of ``scatter_by_slot`` (deepseek_sparse_attention.cpp:200). + + Writes ``value`` rows into the paged ``cache`` at the physical slots given by + ``slot_mapping`` (= block_id * block_size + offset). + """ + if ( + cache is None + or cache.numel() == 0 + or slot_mapping is None + or slot_mapping.numel() == 0 + or value is None + or value.numel() == 0 + ): + return + value_2d = value.reshape(-1, value.size(-1)) + cache_2d = cache.view(-1, value_2d.size(1)) + slots = slot_mapping.reshape(-1).to(torch.long).to(cache.device) + update_rows = min(slots.size(0), value_2d.size(0)) + if update_rows <= 0: + return + valid = slots[:update_rows] >= 0 + if cache.device.type == "npu": + # Match C++ scatter_by_slot exactly. The dedicated NPU kernel preserves + # the cache's storage/layout semantics; Tensor.index_copy_ selects a + # different implementation and left the persistent SWA cache invalid + # for the first decode step. + from xllm.python import kernels + + slots_slice = slots[:update_rows] + safe_slots = slots_slice.clamp_min(0) + valid_mask = slots_slice.ge(0).unsqueeze(1) + old_values = cache_2d.index_select(0, safe_slots) + safe_values = torch.where( + valid_mask, + value_2d[:update_rows].to(cache.dtype), + old_values, + ) + kernels.scatter_nd_update( + cache_2d, + safe_slots.reshape(-1, 1), + safe_values, + ) + return + if not valid.any(): + return + cache_2d.index_copy_( + 0, + slots[:update_rows][valid], + value_2d[:update_rows][valid].to(cache.dtype), + ) + + +def _sparse_attn_sharedkv(**kwargs): + """Thin indirection so the backend can be unit-tested without the kernel.""" + from xllm.python import kernels + + return kernels.sparse_attn_sharedkv(**kwargs) diff --git a/xllm/python/attention/dsa_metadata.py b/xllm/python/attention/dsa_metadata.py new file mode 100644 index 0000000000..86695e9a49 --- /dev/null +++ b/xllm/python/attention/dsa_metadata.py @@ -0,0 +1,665 @@ +# 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. + +"""DeepSeek-V4 DSA metadata construction (faithful Python port of the C++ +``DSAMetadataBuilder`` in ``core/layers/common/dsa_metadata_builder.cpp``). + +This is the Python-path counterpart of the C++ ``build_dsa_fields`` step: it +turns the framework-allocated ``multi_block_tables`` (per-manager block tables, +exposed through ``AttentionMetadataView``) plus the per-layer +``caches_info`` / ``group_infos`` (rebuilt from ``compress_ratios`` + +``window_size``) into the per-layer ``block_tables`` / ``slot_mappings`` the +DSA attention kernel consumes, along with the c4/c128 compressed positions, +sequence-length metadata, and RoPE tables. + +The C++ model forward builds this inside ``DeepseekV4ModelImpl``; under +``--model_impl python`` the C++ forward never runs, so the Python DSA attention +backend builds it here from the same inputs. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Sequence + +import torch + +# --------------------------------------------------------------------------- +# Cache-type enum (mirrors ``DSACacheType`` in dsa_metadata.h). +# --------------------------------------------------------------------------- +DSA_CACHE_TOKEN = 0 +DSA_CACHE_SEQUENCE = 1 +DSA_CACHE_SLIDING_WINDOW = 2 + + +@dataclass +class DSACacheInfo: + """Per-cache descriptor: which group it belongs to and its own shape.""" + + group_id: int + cache_type: int + ratio: int + block_size: int + + +@dataclass +class DSAGroupInfo: + """Per-group descriptor: one block-manager pool.""" + + cache_type: int + ratio: int + block_size: int + + +@dataclass +class DsaMetadata: + """Per-forward DSA metadata, shared across all layers of one model.""" + + # Current layer selected by the model loop. Mirrors C++ + # DSAMetadata::layer_id and is updated immediately before each layer runs. + layer_id: int + + # Sequence lengths (host, int32). + seq_lens: torch.Tensor + seq_lens_q: torch.Tensor + actual_seq_lengths_kv: torch.Tensor + actual_seq_lengths_query: torch.Tensor + kv_cu_seq_lens: torch.Tensor + max_seqlen_kv: torch.Tensor + max_seqlen_q: torch.Tensor + max_query_len: int + max_seq_len: int + + # Positions. + input_positions: torch.Tensor + c4_pad_positions: torch.Tensor + c128_pad_positions: torch.Tensor + start_pos: torch.Tensor + + # RoPE base tables. + cos_table: torch.Tensor | None = None + sin_table: torch.Tensor | None = None + # Per-ratio compressed RoPE tables (C++ DeepseekV4RotaryEmbedding c4/c128 + # groups, compress_rope_theta, no mscale). Populated per-request by indexing + # the compress RoPE cache with c4/c128_pad_positions. + c4_cos: torch.Tensor | None = None + c4_sin: torch.Tensor | None = None + c128_cos: torch.Tensor | None = None + c128_sin: torch.Tensor | None = None + + # block_tables / slot_mappings: [n_layers][n_caches_per_layer]; caches in the + # same group share the same underlying tensor (no copy). + block_tables: list[list[torch.Tensor]] = field(default_factory=list) + slot_mappings: list[list[torch.Tensor]] = field(default_factory=list) + + # Precomputed AICPU tiling metadata (filled by the backend, not the builder). + c1_metadata: torch.Tensor | None = None + c4_metadata: torch.Tensor | None = None + c128_metadata: torch.Tensor | None = None + qli_metadata: torch.Tensor | None = None + hadamard: torch.Tensor | None = None + + # Keep AICPU metadata-builder inputs alive until all asynchronously + # enqueued kernels (and, for ACL graph, the captured graph entry) are done. + # C++ gets this lifetime from the owning DSAMetadata fields; Python creates + # additional empty optional tensors and device copies while precomputing. + precomputed_metadata_inputs: tuple[torch.Tensor, ...] = () + + # Owns the NPU storage for request-shaped metadata packed in prepare(). + # Tensor fields above may be views into this buffer, matching C++ + # DSAMetadata::packed_metadata_buffer. + packed_metadata_buffer: torch.Tensor | None = None + + is_acl_graph: bool = False + + # Per-forward DeepSeek-V4 context-parallel state. This is the Python + # counterpart of DSAMetadata::v4_cp_context; it is populated only for + # prefill when cp_size > 1 and must never survive into a later forward. + v4_cp_context: object | None = None + + +def _normalize_compress_ratio(ratio: int) -> int: + """Mirrors ``deepseek_v4_normalize_compress_ratio``.""" + return 1 if ratio <= 1 else ratio + + +def build_cache_specs( + compress_ratios: Sequence[int], + window_size: int, + n_layers: int, +) -> tuple[list[list[DSACacheInfo]], list[DSAGroupInfo]]: + """Python port of ``deepseek_v4_build_cache_specs`` (deepseek_v4.h:332). + + Builds the per-layer ``caches_info`` and the deduplicated ``group_infos`` + from ``compress_ratios`` + ``window_size``. Group 0 is always the SWA + (sliding-window) group; TOKEN groups for ratios {4, 128} are registered in + the order they first appear. + """ + base_block_size = 128 + group_infos: list[DSAGroupInfo] = [] + group_key_map: dict[tuple[int, int, int], int] = {} + + def register_group(cache_type: int, ratio: int, block_size: int) -> int: + key = (ratio, cache_type, block_size) + gid = group_key_map.get(key) + if gid is not None: + return gid + gid = len(group_infos) + group_key_map[key] = gid + group_infos.append(DSAGroupInfo(cache_type, ratio, block_size)) + return gid + + register_group(DSA_CACHE_SLIDING_WINDOW, 1, window_size) + for ratio in compress_ratios: + cr = _normalize_compress_ratio(ratio) + if cr in (4, 128): + register_group(DSA_CACHE_TOKEN, cr, base_block_size) + + caches_info: list[list[DSACacheInfo]] = [[] for _ in range(n_layers)] + for layer_id in range(n_layers): + cr = compress_ratios[layer_id] if layer_id < len(compress_ratios) else 1 + cr = _normalize_compress_ratio(cr) + + if cr == 1: + entries = [(DSA_CACHE_SLIDING_WINDOW, 1, window_size)] + elif cr == 4: + # cmp_kv, cmp_index, swa, kv_state, score_state, idx_kv, + # idx_score, indexer_scale. + entries = [ + (DSA_CACHE_TOKEN, 4, base_block_size), + (DSA_CACHE_TOKEN, 4, base_block_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_TOKEN, 4, base_block_size), + ] + elif cr == 128: + entries = [ + (DSA_CACHE_TOKEN, 128, base_block_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + (DSA_CACHE_SLIDING_WINDOW, 1, window_size), + ] + else: + entries = [] + + for cache_type, ratio, block_size in entries: + gid = register_group(cache_type, ratio, block_size) + caches_info[layer_id].append(DSACacheInfo(gid, cache_type, ratio, block_size)) + + return caches_info, group_infos + + +class DsaMetadataBuilder: + """Faithful Python port of ``DSAMetadataBuilder`` (dsa_metadata_builder.cpp). + + Construct once per model (with the static ``caches_info`` / ``group_infos``), + then call :meth:`build` every forward to expand the per-manager block tables + into per-layer ``block_tables`` / ``slot_mappings``. + """ + + def __init__( + self, + caches_info: list[list[DSACacheInfo]], + group_infos: list[DSAGroupInfo], + ) -> None: + self.caches_info = caches_info + self.group_infos = group_infos + + # -- public API --------------------------------------------------------- + + def build( + self, + multi_block_tables: Sequence[torch.Tensor], + kv_seq_lens: Sequence[int], + q_seq_lens: Sequence[int] | None, + positions: torch.Tensor, + dsa_cos_sin: torch.Tensor | None, + is_prefill: bool, + is_chunked_prefill: bool, + enable_graph: bool = False, + graph_block_table_capacity_cols: int = 0, + max_query_len: int = 0, + max_seq_len: int = 0, + ) -> DsaMetadata: + batch_size = len(kv_seq_lens) + if q_seq_lens is None or len(q_seq_lens) != batch_size: + if is_prefill or is_chunked_prefill: + q_lens = list(kv_seq_lens) + else: + q_lens = [1] * batch_size + else: + q_lens = list(q_seq_lens) + + dsa = self._build_seq_lengths( + kv_seq_lens, + q_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + ) + dsa.input_positions = positions + dsa.is_acl_graph = enable_graph + if dsa_cos_sin is not None and dsa_cos_sin.numel() > 0: + cos_sin_chunks = dsa_cos_sin.chunk(2, dim=-1) + dsa.cos_table = cos_sin_chunks[0].contiguous() + dsa.sin_table = cos_sin_chunks[1].contiguous() + if positions is not None and positions.numel() > 0: + self._build_positions(dsa, kv_seq_lens, q_lens, enable_graph) + dsa.start_pos = (dsa.actual_seq_lengths_kv - dsa.seq_lens_q).to(torch.int32) + + self._build_block_tables_and_slots( + multi_block_tables, + kv_seq_lens, + q_lens, + batch_size, + positions, + enable_graph, + graph_block_table_capacity_cols, + dsa, + ) + return dsa + + # -- step 1: sequence lengths (build_seq_lengths, cpp:574-661) ---------- + + def _build_seq_lengths( + self, + kv_seq_lens: Sequence[int], + q_lens: Sequence[int], + *, + max_query_len: int, + max_seq_len: int, + ) -> DsaMetadata: + device = torch.device("cpu") + kv = torch.tensor(kv_seq_lens, dtype=torch.int32, device=device) + q = torch.tensor(q_lens, dtype=torch.int32, device=device) + zeros_prefix = torch.zeros(1, dtype=torch.int32, device=device) + actual_seq_lengths_query = torch.cat([zeros_prefix, q.cumsum(0).to(torch.int32)]) + kv_cu = torch.cat([zeros_prefix, kv.cumsum(0).to(torch.int32)]) + max_kv = kv.max().to(torch.int32) if kv.numel() else torch.zeros(1, dtype=torch.int32, device=device) + max_q = q.max().to(torch.int32) if q.numel() else torch.zeros(1, dtype=torch.int32, device=device) + max_query_len = max(int(max_query_len), max((int(value) for value in q_lens), default=0)) + max_seq_len = max( + int(max_seq_len), + max((int(value) for value in kv_seq_lens), default=0), + ) + return DsaMetadata( + layer_id=-1, + seq_lens=kv, + seq_lens_q=q, + actual_seq_lengths_kv=kv, + actual_seq_lengths_query=actual_seq_lengths_query, + kv_cu_seq_lens=kv_cu, + max_seqlen_kv=max_kv, + max_seqlen_q=max_q, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + input_positions=torch.empty(0), + c4_pad_positions=torch.empty(0, dtype=torch.int64), + c128_pad_positions=torch.empty(0, dtype=torch.int64), + start_pos=torch.empty(0), + ) + + # -- step 2: positions (build_positions, cpp:663-777) ------------------ + + def _build_positions( + self, + dsa: DsaMetadata, + kv_seq_lens: Sequence[int], + q_lens: Sequence[int], + enable_graph: bool, + ) -> None: + """Collect c4/c128 compressed RoPE positions. + + For each query token at absolute ``pos``, when ``(pos + 1) % ratio == 0`` + the compressed RoPE needs the position ``next_pos - ratio``. + """ + total_tokens = int(dsa.input_positions.numel()) + c4_positions: list[int] = [] + c128_positions: list[int] = [] + for seq, kv_len in enumerate(kv_seq_lens): + q_len = min(q_lens[seq], kv_len) + start_pos = kv_len - q_len + for i in range(q_len): + pos = start_pos + i + next_pos = pos + 1 + if next_pos % 4 == 0: + c4_positions.append(next_pos - 4) + if next_pos % 128 == 0: + c128_positions.append(next_pos - 128) + + def _pad(positions: list[int], ratio: int) -> torch.Tensor: + if enable_graph: + # Graph mode pads to total_tokens so the tensor address is stable + # across bucket sizes. C++ vector::resize() zero-fills the tail. + out = torch.zeros(total_tokens, dtype=dsa.input_positions.dtype) + for idx, p in enumerate(positions): + out[idx] = p + return out + # Non-graph: resize to min(total_tokens, total_tokens//ratio + batch_size) + # with 0 padding, matching C++ dsa_metadata_builder.cpp:717 + # (c4_target = min(num_tokens, num_tokens/4 + batch_size)). + batch_size = len(kv_seq_lens) + target = min(total_tokens, total_tokens // ratio + batch_size) + out = torch.zeros(target, dtype=dsa.input_positions.dtype) + for idx, p in enumerate(positions): + if idx >= target: + break + out[idx] = p + return out + + dsa.c4_pad_positions = _pad(c4_positions, 4) + dsa.c128_pad_positions = _pad(c128_positions, 128) + + # -- step 3: block_tables / slot_mappings (build_dsa_fields, cpp:169-264) + + def _build_block_tables_and_slots( + self, + multi_block_tables: Sequence[torch.Tensor], + ctx_lens: Sequence[int], + q_lens: Sequence[int], + batch_size: int, + positions: torch.Tensor, + enable_graph: bool, + graph_block_table_capacity_cols: int, + dsa: DsaMetadata, + ) -> None: + if not multi_block_tables or not self.caches_info: + return + + active = list(multi_block_tables) + manager_num = len(active) + # Packed [manager, blocks] auto-unpack when batch_size == 1. + if ( + manager_num == 1 + and batch_size == 1 + and active[0].dim() == 2 + and active[0].size(0) > 1 + and active[0].size(0) <= len(self.group_infos) + ): + packed = active[0].contiguous() + active = [packed[m].unsqueeze(0).contiguous() for m in range(packed.size(0))] + manager_num = len(active) + + if manager_num > len(self.group_infos): + raise ValueError(f"manager count {manager_num} exceeds group count {len(self.group_infos)}") + if enable_graph and graph_block_table_capacity_cols > 0: + for manager_id, block_table in enumerate(active): + if block_table.dim() != 2: + raise ValueError( + f"ACL graph multi_block_tables must be 2-D: manager {manager_id} has rank {block_table.dim()}" + ) + if block_table.size(1) > graph_block_table_capacity_cols: + raise ValueError( + "ACL graph block table exceeds bucket capacity: " + f"manager {manager_id} requires {block_table.size(1)} " + f"columns, capacity is {graph_block_table_capacity_cols}" + ) + + graph_slot_capacity = int(positions.numel()) if enable_graph and positions.numel() > 0 else 0 + total_tokens = sum(int(x) for x in ctx_lens) + + proc_bt: list[torch.Tensor] = [torch.empty(0)] * manager_num + proc_slots: list[torch.Tensor] = [torch.empty(0)] * manager_num + for m in range(manager_num): + gi = self.group_infos[m] + proc_bt[m], proc_slots[m] = self._process_group( + active[m], + gi, + ctx_lens, + q_lens, + batch_size, + total_tokens, + graph_slot_capacity, + graph_block_table_capacity_cols, + ) + + n_layers = len(self.caches_info) + dsa.block_tables = [[] for _ in range(n_layers)] + dsa.slot_mappings = [[] for _ in range(n_layers)] + for lid in range(n_layers): + for ci in range(len(self.caches_info[lid])): + gid = self.caches_info[lid][ci].group_id + if gid < manager_num: + dsa.block_tables[lid].append(proc_bt[gid]) + dsa.slot_mappings[lid].append(proc_slots[gid]) + else: + dsa.block_tables[lid].append(torch.empty(0)) + dsa.slot_mappings[lid].append(torch.empty(0)) + + # -- per-group processing (process_group, cpp:323-362) ----------------- + + def _process_group( + self, + raw_bt: torch.Tensor, + gi: DSAGroupInfo, + ctx_lens: Sequence[int], + q_lens: Sequence[int], + batch_size: int, + total_tokens: int, + graph_slot_capacity: int, + graph_block_table_capacity_cols: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + if gi.cache_type == DSA_CACHE_TOKEN: + return self._process_token_group( + raw_bt, + gi.ratio, + gi.block_size, + ctx_lens, + q_lens, + batch_size, + graph_slot_capacity, + graph_block_table_capacity_cols, + ) + if gi.cache_type == DSA_CACHE_SLIDING_WINDOW: + return self._process_swa_group( + raw_bt, + gi.block_size, + ctx_lens, + q_lens, + batch_size, + graph_slot_capacity, + graph_block_table_capacity_cols, + ) + # SEQUENCE: expand the whole context. + return self._expand_blocks_to_slots(raw_bt, gi, ctx_lens, batch_size, total_tokens) + + # -- TOKEN group (process_token_group, cpp:364-464) -------------------- + # Commits only the compressed rows crossed by the current forward step. + + def _process_token_group( + self, + raw_bt: torch.Tensor, + ratio: int, + block_size: int, + ctx_lens: Sequence[int], + q_lens: Sequence[int], + batch_size: int, + graph_slot_capacity: int, + graph_block_table_capacity_cols: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + committed_rows = 0 + for seq in range(batch_size): + ctx_len = int(ctx_lens[seq]) + q_len = max(0, min(int(q_lens[seq]), ctx_len)) + prev_ctx_len = ctx_len - q_len + committed_rows += ctx_len // ratio - prev_ctx_len // ratio + + out_slot_rows = max(graph_slot_capacity, committed_rows) if graph_slot_capacity > 0 else committed_rows + out_slots = torch.full((out_slot_rows,), -1, dtype=torch.int32, device=raw_bt.device) + semantic_cols = int(raw_bt.size(1)) + + def slot_for_compressed_index(seq: int, compressed_idx: int) -> int: + if seq >= raw_bt.size(0) or semantic_cols <= 0: + return -1 + block_idx = compressed_idx // block_size + if block_idx >= semantic_cols: + return -1 + block_id = int(raw_bt[seq, block_idx].item()) + if block_id < 0: + return -1 + block_offset = compressed_idx % block_size + return block_id * block_size + block_offset + + write_idx = 0 + slots_list = out_slots.tolist() + for seq in range(batch_size): + ctx_len = int(ctx_lens[seq]) + q_len = max(0, min(int(q_lens[seq]), ctx_len)) + prev_ctx_len = ctx_len - q_len + prev_committed = prev_ctx_len // ratio + committed = ctx_len // ratio + new_committed = committed - prev_committed + for i in range(new_committed): + slots_list[write_idx] = slot_for_compressed_index(seq, prev_committed + i) + write_idx += 1 + out_slots = torch.tensor(slots_list, dtype=torch.int32, device=raw_bt.device) + + out_bt = raw_bt + if graph_slot_capacity > 0 and graph_block_table_capacity_cols > 0: + cap = max(graph_block_table_capacity_cols, int(raw_bt.size(1))) + out_bt = self._pad_block_table(raw_bt, batch_size, cap, -1) + return out_bt, out_slots + + # -- SWA group (process_swa_group, cpp:466-572) ------------------------ + # Writes only the current forward's query tokens, ring-indexed by position. + + def _process_swa_group( + self, + raw_bt: torch.Tensor, + block_size: int, + ctx_lens: Sequence[int], + q_lens: Sequence[int], + batch_size: int, + graph_slot_capacity: int, + graph_block_table_capacity_cols: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + query_total_tokens = 0 + for seq in range(batch_size): + query_total_tokens += max(0, min(int(q_lens[seq]), int(ctx_lens[seq]))) + + out_slot_rows = max(graph_slot_capacity, query_total_tokens) if graph_slot_capacity > 0 else query_total_tokens + out_slots = torch.full((out_slot_rows,), -1, dtype=torch.int32, device=raw_bt.device) + semantic_cols = int(raw_bt.size(1)) + + def slot_for_position(seq: int, pos: int) -> int: + if semantic_cols <= 0 or seq >= raw_bt.size(0): + return -1 + block_idx = (pos // block_size) % semantic_cols + block_id = int(raw_bt[seq, block_idx].item()) + if block_id < 0: + return -1 + block_offset = pos % block_size + return block_id * block_size + block_offset + + write_idx = 0 + slots_list = out_slots.tolist() + for seq in range(batch_size): + ctx_len = int(ctx_lens[seq]) + q_len = max(0, min(int(q_lens[seq]), ctx_len)) + if seq >= raw_bt.size(0): + write_idx += q_len + continue + q_start = ctx_len - q_len + for i in range(q_len): + slots_list[write_idx] = slot_for_position(seq, q_start + i) + write_idx += 1 + out_slots = torch.tensor(slots_list, dtype=torch.int32, device=raw_bt.device) + + # Rebuild the read-side block table: keep only the SWA window columns, + # right-aligned. + dst_lens = [(max(int(ctx_lens[s]), 0) + block_size - 1) // block_size for s in range(batch_size)] + max_dst_len = max(max(dst_lens) if dst_lens else 0, semantic_cols) + if graph_slot_capacity > 0 and graph_block_table_capacity_cols > 0: + storage_cols = max(graph_block_table_capacity_cols, int(raw_bt.size(1))) + max_dst_len = max(max_dst_len, storage_cols) + new_bt = torch.full( + (batch_size, max_dst_len), + -1, + dtype=torch.int32, + device=raw_bt.device, + ) + for s in range(batch_size): + if s >= raw_bt.size(0): + continue + retained_cols = min(semantic_cols, dst_lens[s]) + start_col = dst_lens[s] - retained_cols + for j in range(retained_cols): + logical_col = start_col + j + physical_col = logical_col % semantic_cols + new_bt[s, logical_col] = raw_bt[s, physical_col] + return new_bt, out_slots + + # -- SEQUENCE group (expand_blocks_to_slots, cpp:270-307) -------------- + + def _expand_blocks_to_slots( + self, + raw_bt: torch.Tensor, + gi: DSAGroupInfo, + ctx_lens: Sequence[int], + batch_size: int, + total_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + block_size = gi.block_size + slots = torch.full((total_tokens,), -1, dtype=torch.int32, device=raw_bt.device) + max_blocks = int(raw_bt.size(1)) + start_idx = 0 + for seq in range(batch_size): + token_len = int(ctx_lens[seq]) + slot_num = self._compute_slot_num(gi, token_len) + if seq >= raw_bt.size(0): + start_idx += token_len + continue + filled = 0 + for blk in range(max_blocks): + if filled >= slot_num: + break + block_id = int(raw_bt[seq, blk].item()) + if block_id < 0: + break + for off in range(block_size): + if filled >= slot_num: + break + slots[start_idx + filled] = block_id * block_size + off + filled += 1 + start_idx += token_len + # Replace -1 padding with 0 (C++ does torch::where(eq(-1), 0, raw)). + slots = torch.where(slots.eq(-1), torch.zeros_like(slots), slots) + return raw_bt, slots + + @staticmethod + def _compute_slot_num(gi: DSAGroupInfo, token_len: int) -> int: + if gi.cache_type == DSA_CACHE_TOKEN: + return token_len // gi.ratio + # SLIDING_WINDOW + block_size = gi.block_size + if token_len > block_size: + return token_len % block_size + block_size + remainder = token_len % block_size + return block_size if (remainder == 0 and token_len > 0) else remainder + + @staticmethod + def _pad_block_table( + raw_bt: torch.Tensor, + batch_size: int, + capacity_cols: int, + pad_value: int, + ) -> torch.Tensor: + cols = max(capacity_cols, int(raw_bt.size(1))) + out = torch.full((batch_size, cols), pad_value, dtype=torch.int32, device=raw_bt.device) + rows = min(batch_size, int(raw_bt.size(0))) + copy_cols = min(int(raw_bt.size(1)), cols) + out[:rows, :copy_cols] = raw_bt[:rows, :copy_cols] + return out diff --git a/xllm/python/distributed/__init__.py b/xllm/python/distributed/__init__.py index b0da8b63b9..1795d6864b 100644 --- a/xllm/python/distributed/__init__.py +++ b/xllm/python/distributed/__init__.py @@ -24,6 +24,10 @@ cp_world_size, init_process_group, init_tp_group, + moe_ep_all_reduce, + moe_tp_all_reduce, + tp_all_gather, + tp_all_reduce, tp_rank, ) @@ -33,6 +37,10 @@ "tp_rank", "cp_rank", "cp_world_size", + "tp_all_reduce", + "tp_all_gather", + "moe_tp_all_reduce", + "moe_ep_all_reduce", "all_reduce_", "all_gather", "all_gather_variable", diff --git a/xllm/python/distributed/collectives.py b/xllm/python/distributed/collectives.py index dd7358df11..8785feb376 100644 --- a/xllm/python/distributed/collectives.py +++ b/xllm/python/distributed/collectives.py @@ -57,9 +57,7 @@ def _backend_for(device: torch.device) -> str: return "hccl" -def _shared_store( - host: str, port: int, global_rank: int, global_world_size: int -) -> dist.Store: +def _shared_store(host: str, port: int, global_rank: int, global_world_size: int) -> dist.Store: store_key = (host, port) store = _stores.get(store_key) if store is None: @@ -88,10 +86,7 @@ def _exchange_world_topology( local = {"hostname": socket.gethostname(), "device_index": device_index} key_prefix = "xllm/python_collectives/topology/v1" store.set(f"{key_prefix}/{global_rank}", json.dumps(local)) - return [ - json.loads(store.get(f"{key_prefix}/{rank}").decode("utf-8")) - for rank in range(global_world_size) - ] + return [json.loads(store.get(f"{key_prefix}/{rank}").decode("utf-8")) for rank in range(global_world_size)] def _ensure_world( @@ -118,35 +113,22 @@ def _ensure_world( world_size=global_world_size, timeout=timedelta(minutes=5), ) - _world_topology = _exchange_world_topology( - store, device, global_rank, global_world_size - ) + _world_topology = _exchange_world_topology(store, device, global_rank, global_world_size) _world_initialized = True -def _group_memberships( - group_name: str, world_size: int, global_world_size: int -) -> list[list[int]]: +def _group_memberships(group_name: str, world_size: int, global_world_size: int) -> list[list[int]]: """Every group of this kind, in an order all ranks agree on. ``new_group`` is collective over the whole world, so each rank has to create all groups of a kind in the same order, not only the one it belongs to. """ if world_size <= 0 or global_world_size % world_size: - raise ValueError( - f"{group_name} size {world_size} does not divide the world size " - f"{global_world_size}" - ) + raise ValueError(f"{group_name} size {world_size} does not divide the world size {global_world_size}") count = global_world_size // world_size if group_name in _CONTIGUOUS_GROUPS: - return [ - [index * world_size + offset for offset in range(world_size)] - for index in range(count) - ] - return [ - [index + offset * count for offset in range(world_size)] - for index in range(count) - ] + return [[index * world_size + offset for offset in range(world_size)] for index in range(count)] + return [[index + offset * count for offset in range(world_size)] for index in range(count)] def _supports_symmetric_memory(device: torch.device, ranks: list[int]) -> bool: @@ -201,9 +183,7 @@ def init_process_group( own_ranks = None memberships = _group_memberships(group_name, world_size, global_world_size) for index, ranks in enumerate(memberships): - candidate = dist.new_group( - ranks=ranks, timeout=timedelta(minutes=5), backend=backend - ) + candidate = dist.new_group(ranks=ranks, timeout=timedelta(minutes=5), backend=backend) if global_rank in ranks: own = candidate own_index = index @@ -215,10 +195,7 @@ def init_process_group( f"{world_size}, got {None if own is None else own.rank()}" ) if own_index != group_index: - raise RuntimeError( - f"derived {group_name} group index {own_index} does not match the " - f"caller's {group_index}" - ) + raise RuntimeError(f"derived {group_name} group index {own_index} does not match the caller's {group_index}") assert own_ranks is not None _groups[group_key] = own @@ -252,10 +229,7 @@ def init_tp_group( def _require_group(x: torch.Tensor, group_name: str) -> ProcessGroup: group = _groups.get((group_name, str(x.device))) if group is None: - raise RuntimeError( - f"{group_name} collective called before its process group was " - f"initialized for {x.device}" - ) + raise RuntimeError(f"{group_name} collective called before its process group was initialized for {x.device}") return group @@ -281,6 +255,46 @@ def cp_world_size(device: torch.device | str) -> int: return group.size() if group is not None else 1 +def _native_runtime_op(name: str): + """Return an embedded C++ collective when running under PyExecutorImpl.""" + try: + import xllm_runtime + except ImportError: + return None + return getattr(xllm_runtime, name, None) + + +def tp_all_reduce(x: torch.Tensor) -> None: + op = _native_runtime_op("tp_all_reduce") + if op is not None: + op(x) + return + all_reduce_(x, "tp") + + +def tp_all_gather(x: torch.Tensor, dim: int, world_size: int) -> torch.Tensor: + op = _native_runtime_op("tp_all_gather") + if op is not None: + return op(x, dim) + return all_gather(x, dim, world_size, "tp") + + +def moe_tp_all_reduce(x: torch.Tensor) -> None: + op = _native_runtime_op("moe_tp_all_reduce") + if op is not None: + op(x) + return + all_reduce_(x, "moe_tp") + + +def moe_ep_all_reduce(x: torch.Tensor) -> None: + op = _native_runtime_op("moe_ep_all_reduce") + if op is not None: + op(x) + return + all_reduce_(x, "moe_ep") + + # A one-shot symmetric-memory reduction is an ordinary kernel on the current # stream, so a captured graph runs it inline. NCCL runs collectives on its own # stream, which costs a fork/join per call -- measured at ~32us of device idle @@ -295,9 +309,7 @@ def cp_world_size(device: torch.device | str) -> int: _SYMM_MEM_DTYPES = frozenset((torch.float32, torch.bfloat16)) -def _symm_buffer( - group: ProcessGroup, group_name: str, x: torch.Tensor -) -> torch.Tensor | None: +def _symm_buffer(group: ProcessGroup, group_name: str, x: torch.Tensor) -> torch.Tensor | None: """Return a symmetric-memory staging buffer for ``x``, or None. Allocation and rendezvous are collective and cannot run inside a graph @@ -337,9 +349,7 @@ def all_reduce_(x: torch.Tensor, group_name: str = "tp") -> None: return flat = x.view(-1) buffer.copy_(flat) - torch.ops.symm_mem.one_shot_all_reduce_out( - buffer, "sum", group.group_name, flat - ) + torch.ops.symm_mem.one_shot_all_reduce_out(buffer, "sum", group.group_name, flat) @all_reduce_.register_fake @@ -348,24 +358,17 @@ def _(x: torch.Tensor, group_name: str = "tp") -> None: @torch.library.custom_op("xllm_ops::all_gather", mutates_args=()) -def all_gather( - x: torch.Tensor, dim: int, world_size: int, group_name: str = "tp" -) -> torch.Tensor: +def all_gather(x: torch.Tensor, dim: int, world_size: int, group_name: str = "tp") -> torch.Tensor: group = _require_group(x, group_name) if group.size() != world_size: - raise RuntimeError( - f"{group_name} world-size mismatch: expected {world_size}, " - f"got {group.size()}" - ) + raise RuntimeError(f"{group_name} world-size mismatch: expected {world_size}, got {group.size()}") chunks = [torch.empty_like(x) for _ in range(world_size)] dist.all_gather(chunks, x, group=group) return torch.cat(chunks, dim=dim) @all_gather.register_fake -def _( - x: torch.Tensor, dim: int, world_size: int, group_name: str = "tp" -) -> torch.Tensor: +def _(x: torch.Tensor, dim: int, world_size: int, group_name: str = "tp") -> torch.Tensor: shape = list(x.shape) shape[dim] *= world_size return x.new_empty(shape) @@ -389,9 +392,7 @@ def all_gather_variable( local_tokens = token_counts[rank] if local_tokens < 0 or local_tokens > x.shape[0]: - raise RuntimeError( - f"invalid local token count {local_tokens} for input with {x.shape[0]} rows" - ) + raise RuntimeError(f"invalid local token count {local_tokens} for input with {x.shape[0]} rows") padded_tokens = max(max(token_counts, default=0), 1) padded_shape = list(x.shape) padded_shape[0] = padded_tokens @@ -401,9 +402,7 @@ def all_gather_variable( chunks = [torch.empty_like(padded) for _ in token_counts] dist.all_gather(chunks, padded, group=group) - valid_chunks = [ - chunk[:count] for chunk, count in zip(chunks, token_counts) if count - ] + valid_chunks = [chunk[:count] for chunk, count in zip(chunks, token_counts) if count] if not valid_chunks: empty_shape = list(x.shape) empty_shape[0] = 0 @@ -430,6 +429,10 @@ def _( "tp_rank", "cp_rank", "cp_world_size", + "tp_all_reduce", + "tp_all_gather", + "moe_tp_all_reduce", + "moe_ep_all_reduce", "all_reduce_", "all_gather", "all_gather_variable", diff --git a/xllm/python/kernels_cuda/__init__.py b/xllm/python/kernels_cuda/__init__.py index 32033e0785..613a2c07de 100644 --- a/xllm/python/kernels_cuda/__init__.py +++ b/xllm/python/kernels_cuda/__init__.py @@ -54,6 +54,7 @@ cutlass_fused_moe, fused_moe, grouped_moe, + grouped_moe_with_selected_experts, moe_fused_topk, prepare_grouped_moe_weights, supports_cutlass_moe, @@ -100,6 +101,7 @@ "cutlass_fused_moe", "fused_moe", "grouped_moe", + "grouped_moe_with_selected_experts", "prepare_grouped_moe_weights", "supports_cutlass_moe", "prepare_row_parallel_weight", diff --git a/xllm/python/kernels_cuda/moe.py b/xllm/python/kernels_cuda/moe.py index fa742f4d74..91f0279fa1 100644 --- a/xllm/python/kernels_cuda/moe.py +++ b/xllm/python/kernels_cuda/moe.py @@ -56,9 +56,7 @@ def moe_fused_topk( Returns: Routing weights and expert indices, both ``[num_tokens, topk]``. """ - return torch.ops.xllm_ops.moe_fused_topk( - gating_output, topk, renormalize, scoring_func - ) + return torch.ops.xllm_ops.moe_fused_topk(gating_output, topk, renormalize, scoring_func) def cutlass_fused_moe( @@ -199,9 +197,42 @@ def grouped_moe( active_expert_range, ) raise NotImplementedError( - "grouped_moe has no CUDA kernel; the equivalent CUDA path is " - "moe_fused_topk followed by cutlass_fused_moe" + "grouped_moe has no CUDA kernel; the equivalent CUDA path is moe_fused_topk followed by cutlass_fused_moe" + ) + + +def grouped_moe_with_selected_experts( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + """Reject the NPU-specific pre-selected grouped MoE contract on CUDA.""" + del ( + hidden_states, + topk_weights, + topk_ids, + w13, + w2, + w13_scale, + w2_scale, + w13_offset, + w2_offset, + num_total_experts, + start_expert_id, + num_experts_per_rank, + swiglu_limit, ) + raise NotImplementedError("grouped_moe_with_selected_experts is currently implemented only for NPU") __all__ = [ @@ -211,4 +242,5 @@ def grouped_moe( "fused_moe", "prepare_grouped_moe_weights", "grouped_moe", + "grouped_moe_with_selected_experts", ] diff --git a/xllm/python/kernels_npu/__init__.py b/xllm/python/kernels_npu/__init__.py index dffb94fe5d..a1f7ee519a 100644 --- a/xllm/python/kernels_npu/__init__.py +++ b/xllm/python/kernels_npu/__init__.py @@ -42,6 +42,17 @@ causal_conv1d_decode, causal_conv1d_prefill, ) +from .dsa import ( + compressor, + dequant_swiglu_quant, + hc_post, + hc_pre, + moe_gating_top_k_hash, + quant_lightning_indexer, + quant_lightning_indexer_metadata, + sparse_attn_sharedkv, + sparse_attn_sharedkv_metadata, +) from .gated_delta_net import ( chunk_gated_delta_rule, fused_gdn_prefill_post_conv, @@ -53,6 +64,7 @@ cutlass_fused_moe, fused_moe, grouped_moe, + grouped_moe_with_selected_experts, moe_fused_topk, prepare_grouped_moe_weights, supports_cutlass_moe, @@ -61,6 +73,7 @@ fused_add_rms_norm, l2_norm, rms_norm, + rms_norm_dynamic_quant, rms_norm_gated, ) from .quantization import ( @@ -72,6 +85,7 @@ fused_qk_norm_rope, interleaved_rotary_embedding, mrope, + npu_inplace_partial_rotary_mul, vision_rotary_mul, ) from .sparse_attention import ( @@ -85,6 +99,7 @@ __all__ = [ "rms_norm", "fused_add_rms_norm", + "rms_norm_dynamic_quant", "l2_norm", "rms_norm_gated", "silu_and_mul", @@ -93,12 +108,14 @@ "vision_fusion_attention", "fused_qk_norm_rope", "interleaved_rotary_embedding", + "npu_inplace_partial_rotary_mul", "mrope", "vision_rotary_mul", "moe_fused_topk", "cutlass_fused_moe", "fused_moe", "grouped_moe", + "grouped_moe_with_selected_experts", "prepare_grouped_moe_weights", "supports_cutlass_moe", "prepare_row_parallel_weight", @@ -112,6 +129,15 @@ "sparse_flash_attention_out", "causal_conv1d_prefill", "causal_conv1d_decode", + "compressor", + "dequant_swiglu_quant", + "hc_pre", + "hc_post", + "moe_gating_top_k_hash", + "quant_lightning_indexer", + "quant_lightning_indexer_metadata", + "sparse_attn_sharedkv", + "sparse_attn_sharedkv_metadata", "resolve_gdn_prefill_backend", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", diff --git a/xllm/python/kernels_npu/_custom_op.py b/xllm/python/kernels_npu/_custom_op.py index d575d9a308..4efe748a0a 100644 --- a/xllm/python/kernels_npu/_custom_op.py +++ b/xllm/python/kernels_npu/_custom_op.py @@ -168,9 +168,7 @@ def _dynamic_quant_fake( del smooth_scales, group_index if dst_type == torch.quint4x2: if input.shape[-1] % 8: - raise ValueError( - "dynamic_quant int4 input's last dimension must be divisible by 8" - ) + raise ValueError("dynamic_quant int4 input's last dimension must be divisible by 8") output_shape = (*input.shape[:-1], input.shape[-1] // 8) output_dtype = torch.int32 else: @@ -181,6 +179,22 @@ def _dynamic_quant_fake( return output, scale +def _group_gemm_fake( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor | None, + per_token_scale: torch.Tensor | None, + group_list: torch.Tensor, + split_item: int, + group_type: int, + group_list_type: int, + output_dtype: torch.dtype | None, +) -> torch.Tensor: + del scale, per_token_scale, group_list, split_item, group_type, group_list_type + dtype = output_dtype if output_dtype is not None else x.dtype + return x.new_empty((x.size(0), weight.size(-1)), dtype=dtype) + + def _lightning_indexer_fake( query: torch.Tensor, key: torch.Tensor, @@ -328,20 +342,407 @@ def _sparse_flash_attention_out_fake( return output +# --------------------------------------------------------------------------- +# DeepSeek-V4 DSA kernel fakes +# --------------------------------------------------------------------------- + +# Matches kDsaMetadataBufferElements in xllm_ops_api.h. +_DSA_METADATA_BUFFER_ELEMENTS = 1024 + + +def _rms_norm_dynamic_quant_fake( + input: torch.Tensor, weight: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + del weight, eps + return input.new_empty(input.shape, dtype=torch.int8), input.new_empty(input.shape[:-1], dtype=torch.float32) + + +def _npu_inplace_partial_rotary_mul_fake( + x: torch.Tensor, + r1: torch.Tensor, + r2: torch.Tensor, + rotary_mode: str, + partial_slice: list[int], +) -> None: + del x, r1, r2, rotary_mode, partial_slice + + +def _hc_pre_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + hc_sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del hc_fn, hc_scale, hc_base, hc_sinkhorn_iters, norm_eps, hc_eps + if x.dim() == 4: + y_shape = (x.size(0), x.size(1), x.size(3)) + post_shape = (x.size(0), x.size(1), hc_mult) + comb_shape = (x.size(0), x.size(1), hc_mult, hc_mult) + else: + y_shape = (x.size(0), x.size(2)) + post_shape = (x.size(0), hc_mult) + comb_shape = (x.size(0), hc_mult, hc_mult) + attn_input = x.new_empty(y_shape, dtype=x.dtype) + post = x.new_empty(post_shape, dtype=torch.float32) + comb = x.new_empty(comb_shape, dtype=torch.float32) + return attn_input, post, comb + + +def _hc_post_fake( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + del post, comb + # hc_post returns [T, hc_mult, hidden] (the merged residual streams). + return residual.new_empty(residual.shape, dtype=residual.dtype) + + +def _compressor_fake( + x: torch.Tensor, + wkv: torch.Tensor, + wgate: torch.Tensor, + kv_state: torch.Tensor, + score_state: torch.Tensor, + ape: torch.Tensor, + norm_weight: torch.Tensor, + rope_sin: torch.Tensor, + rope_cos: torch.Tensor, + kv_block_table: torch.Tensor | None, + score_block_table: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + seqused: torch.Tensor | None, + start_pos: torch.Tensor | None, + rope_head_dim: int, + cmp_ratio: int, + coff: int, + norm_eps: float, + rotary_mode: int, + enable_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del ( + wkv, + wgate, + kv_state, + score_state, + ape, + rope_cos, + kv_block_table, + score_block_table, + cu_seqlens, + seqused, + start_pos, + rope_head_dim, + norm_eps, + rotary_mode, + ) + head_dim = norm_weight.size(0) + if x.dim() == 3: + compressed_seq = (x.size(1) + cmp_ratio - 1) // cmp_ratio + cmp_kv_shape = (x.size(0), compressed_seq, head_dim) + grad_shapes = ( + (x.size(0), x.size(1), coff * head_dim), + (x.size(0), compressed_seq, coff * cmp_ratio, head_dim), + (x.size(0), compressed_seq, head_dim), + (x.size(0), compressed_seq), + ) + else: + compressed_seq = rope_sin.size(0) + cmp_kv_shape = (compressed_seq, head_dim) + grad_shapes = ( + (x.size(0), coff * head_dim), + (compressed_seq, coff * cmp_ratio, head_dim), + (compressed_seq, head_dim), + (compressed_seq,), + ) + outputs = [x.new_empty(cmp_kv_shape, dtype=x.dtype)] + if enable_grad: + outputs.extend(x.new_empty(shape, dtype=x.dtype) for shape in grad_shapes) + else: + outputs.extend(x.new_empty((0,), dtype=x.dtype) for _ in grad_shapes) + return tuple(outputs) # type: ignore[return-value] + + +def _sparse_attn_sharedkv_fake( + q: torch.Tensor, + ori_kv: torch.Tensor | None, + cmp_kv: torch.Tensor | None, + ori_sparse_indices: torch.Tensor | None, + cmp_sparse_indices: torch.Tensor | None, + ori_block_table: torch.Tensor | None, + cmp_block_table: torch.Tensor | None, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + sinks: torch.Tensor | None, + metadata: torch.Tensor | None, + softmax_scale: float, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + return_softmax_lse: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + ori_kv, + cmp_kv, + ori_sparse_indices, + cmp_sparse_indices, + ori_block_table, + cmp_block_table, + sinks, + metadata, + softmax_scale, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + return_softmax_lse, + ) + out = q.new_empty(q.shape, dtype=q.dtype) + lse = q.new_empty((0,), dtype=q.dtype) + return out, lse + + +def _sparse_attn_sharedkv_metadata_fake( + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_kv: int, + ori_topk: int, + cmp_topk: int, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + has_ori_kv: bool, + has_cmp_kv: bool, +) -> torch.Tensor: + del ( + num_heads_q, + num_heads_kv, + head_dim, + batch_size, + max_seqlen_q, + max_seqlen_kv, + ori_topk, + cmp_topk, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + has_ori_kv, + has_cmp_kv, + ) + for tensor in ( + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + ): + if tensor is not None: + return tensor.new_empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + return torch.empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + + +def _quant_lightning_indexer_fake( + query: torch.Tensor, + key: torch.Tensor, + weights: torch.Tensor, + query_dequant_scale: torch.Tensor, + key_dequant_scale: torch.Tensor, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + block_table: torch.Tensor | None, + metadata: torch.Tensor | None, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + return_value: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + weights, + query_dequant_scale, + key_dequant_scale, + query_quant_mode, + key_quant_mode, + block_table, + metadata, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + ) + key_head_num = key.size(1) if layout_key == "TND" else key.size(2) + if layout_query == "BSND": + out_shape = (query.size(0), query.size(1), key_head_num, sparse_count) + else: + out_shape = (query.size(0), key_head_num, sparse_count) + out = query.new_zeros(out_shape, dtype=torch.int32) + val = ( + query.new_empty(out_shape, dtype=torch.float32) if return_value else query.new_empty((0,), dtype=torch.float32) + ) + return out, val + + +def _quant_lightning_indexer_metadata_fake( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_k: int, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + device: str, +) -> torch.Tensor: + del ( + num_heads_q, + num_heads_k, + head_dim, + query_quant_mode, + key_quant_mode, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + device, + ) + for tensor in (actual_seq_lengths_query, actual_seq_lengths_key): + if tensor is not None: + return tensor.new_empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + return torch.empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + + register_fake("xllm_ops::rms_norm", _rms_norm_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) -register_fake( - "xllm_ops::update_decode_graph_metadata", _update_decode_graph_metadata_fake -) +register_fake("xllm_ops::update_decode_graph_metadata", _update_decode_graph_metadata_fake) register_fake("xllm_ops::quant_matmul", _quant_matmul_fake) register_fake("xllm_ops::quantize_per_tensor", _quantize_per_tensor_fake) register_fake("xllm_ops::dynamic_quant", _dynamic_quant_fake) +register_fake("xllm_ops::group_gemm", _group_gemm_fake) register_fake("xllm_ops::lightning_indexer", _lightning_indexer_fake) register_fake("xllm_ops::lightning_indexer_out", _lightning_indexer_out_fake) register_fake("xllm_ops::scatter_nd_update", _scatter_nd_update_fake) register_fake("xllm_ops::sparse_flash_attention", _sparse_flash_attention_fake) +register_fake("xllm_ops::sparse_flash_attention_out", _sparse_flash_attention_out_fake) +register_fake("xllm_ops::rms_norm_dynamic_quant", _rms_norm_dynamic_quant_fake) +register_fake( + "xllm_ops::npu_inplace_partial_rotary_mul", + _npu_inplace_partial_rotary_mul_fake, +) +register_fake("xllm_ops::compressor", _compressor_fake) + + +def _moe_gating_top_k_hash_fake( + x: torch.Tensor, + k: int, + bias: torch.Tensor | None, + input_ids: torch.Tensor | None, + tid2eid: torch.Tensor | None, + k_group: int, + group_count: int, + routed_scaling_factor: float, + eps: float, + group_select_mode: int, + renorm: int, + norm_type: int, + out_flag: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del bias, input_ids, tid2eid, k_group, group_count, routed_scaling_factor + del eps, group_select_mode, renorm, norm_type, out_flag + y_shape = (*x.shape[:-1], k) + y = x.new_empty(y_shape, dtype=x.dtype) + expert_idx = x.new_empty(y_shape, dtype=torch.int32) + out = x.new_empty(x.shape, dtype=torch.float32) + return y, expert_idx, out + + +register_fake("xllm_ops::moe_gating_top_k_hash", _moe_gating_top_k_hash_fake) + + +def _dequant_swiglu_quant_fake( + x: torch.Tensor, + weight_scale: torch.Tensor | None, + activation_scale: torch.Tensor | None, + bias: torch.Tensor | None, + quant_scale: torch.Tensor | None, + quant_offset: torch.Tensor | None, + group_index: torch.Tensor | None, + activate_left: bool, + quant_mode: int, + swiglu_mode: int, + clamp_limit: float, + glu_alpha: float, + glu_bias: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del weight_scale, activation_scale, bias, quant_scale, quant_offset + del group_index, activate_left, quant_mode, swiglu_mode + del clamp_limit, glu_alpha, glu_bias + # Output is half of input's last dim (SwiGLU splits gate/up). + out_dim = x.size(-1) // 2 + act_quantized = x.new_empty(x.size(0), out_dim, dtype=torch.int8) + act_scale = x.new_empty(x.shape[:-1], dtype=torch.float32) + return act_quantized, act_scale + + +register_fake("xllm_ops::dequant_swiglu_quant", _dequant_swiglu_quant_fake) +register_fake("xllm_ops::hc_pre", _hc_pre_fake) +register_fake("xllm_ops::hc_post", _hc_post_fake) +register_fake("xllm_ops::sparse_attn_sharedkv", _sparse_attn_sharedkv_fake) +register_fake("xllm_ops::sparse_attn_sharedkv_metadata", _sparse_attn_sharedkv_metadata_fake) +register_fake("xllm_ops::quant_lightning_indexer", _quant_lightning_indexer_fake) register_fake( - "xllm_ops::sparse_flash_attention_out", _sparse_flash_attention_out_fake + "xllm_ops::quant_lightning_indexer_metadata", + _quant_lightning_indexer_metadata_fake, ) diff --git a/xllm/python/kernels_npu/dsa.py b/xllm/python/kernels_npu/dsa.py new file mode 100644 index 0000000000..0382977c26 --- /dev/null +++ b/xllm/python/kernels_npu/dsa.py @@ -0,0 +1,372 @@ +# 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. + +"""NPU DeepSeek-V4 DSA kernels. + +These wrap the AscendC operators registered as ``torch.ops.xllm_ops.*`` by +``core/kernels/npu/npu_ops_library.cpp``. They drive the two-stage sparse +attention (original + compressed KV), the KV compressor, the quantized +lightning indexer, and the HyperConnection pre/post used by DeepSeek-V4's DSA +attention path. +""" + +from __future__ import annotations + +import torch + + +def dequant_swiglu_quant( + x: torch.Tensor, + weight_scale: torch.Tensor | None, + activation_scale: torch.Tensor | None, + bias: torch.Tensor | None = None, + quant_scale: torch.Tensor | None = None, + quant_offset: torch.Tensor | None = None, + group_index: torch.Tensor | None = None, + activate_left: bool = True, + quant_mode: int = 1, + swiglu_mode: int = 1, + clamp_limit: float = 0.0, + glu_alpha: float = 1.0, + glu_bias: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused dequant + SwiGLU + dynamic quant (replaces manual loop).""" + return torch.ops.xllm_ops.dequant_swiglu_quant( + x, + weight_scale, + activation_scale, + bias, + quant_scale, + quant_offset, + group_index, + activate_left, + quant_mode, + swiglu_mode, + clamp_limit, + glu_alpha, + glu_bias, + ) + + +def moe_gating_top_k_hash( + x: torch.Tensor, + k: int, + bias: torch.Tensor | None, + input_ids: torch.Tensor | None, + tid2eid: torch.Tensor | None, + k_group: int, + group_count: int, + routed_scaling_factor: float, + eps: float, + group_select_mode: int, + renorm: int, + norm_type: int, + out_flag: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """DeepSeek-V4 MoE hash routing gate.""" + return torch.ops.xllm_ops.moe_gating_top_k_hash( + x, + k, + bias, + input_ids, + tid2eid, + k_group, + group_count, + routed_scaling_factor, + eps, + group_select_mode, + renorm, + norm_type, + out_flag, + ) + + +def hc_pre( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + hc_sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """HyperConnection pre: mix hc_mult streams into one sub-block input. + + Returns ``(attn_input, post, comb)`` where post/comb feed ``hc_post``. + """ + return torch.ops.xllm_ops.hc_pre(x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps) + + +def hc_post( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + """HyperConnection post: combine sub-block output with the residual streams.""" + return torch.ops.xllm_ops.hc_post(x, residual, post, comb) + + +def compressor( + x: torch.Tensor, + wkv: torch.Tensor, + wgate: torch.Tensor, + kv_state: torch.Tensor, + score_state: torch.Tensor, + ape: torch.Tensor, + norm_weight: torch.Tensor, + rope_sin: torch.Tensor, + rope_cos: torch.Tensor, + kv_block_table: torch.Tensor | None, + score_block_table: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + seqused: torch.Tensor | None, + start_pos: torch.Tensor | None, + rope_head_dim: int, + cmp_ratio: int, + coff: int, + norm_eps: float, + rotary_mode: int, + enable_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pool KV along the token axis by ``cmp_ratio`` (NSA-style compressor). + + ``kv_state`` and ``score_state`` are updated in place. + + Returns ``(cmp_kv, wkv_proj, softmax_res, norm_x, norm_rstd)``; only + ``cmp_kv`` is consumed by the DSA path. + """ + # C++ moves DSA metadata to the active device before dispatch. Keep this + # adapter deterministic; experimental clone/noalias paths do not belong in + # the public binding. + kv_block_table = kv_block_table.to(x.device) if kv_block_table is not None else None + score_block_table = score_block_table.to(x.device) if score_block_table is not None else None + cu_seqlens = cu_seqlens.to(x.device) if cu_seqlens is not None else None + seqused = seqused.to(x.device) if seqused is not None else None + start_pos = start_pos.to(x.device) if start_pos is not None else None + return torch.ops.xllm_ops.compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table, + score_block_table, + cu_seqlens, + seqused, + start_pos, + rope_head_dim, + cmp_ratio, + coff, + norm_eps, + rotary_mode, + enable_grad, + ) + + +def sparse_attn_sharedkv( + q: torch.Tensor, + ori_kv: torch.Tensor | None, + cmp_kv: torch.Tensor | None, + ori_sparse_indices: torch.Tensor | None, + cmp_sparse_indices: torch.Tensor | None, + ori_block_table: torch.Tensor | None, + cmp_block_table: torch.Tensor | None, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + sinks: torch.Tensor | None, + metadata: torch.Tensor | None, + softmax_scale: float, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + return_softmax_lse: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Two-stage sparse attention over original and compressed KV.""" + return torch.ops.xllm_ops.sparse_attn_sharedkv( + q, + ori_kv, + cmp_kv, + ori_sparse_indices, + cmp_sparse_indices, + ori_block_table, + cmp_block_table, + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + sinks, + metadata, + softmax_scale, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + return_softmax_lse, + ) + + +def sparse_attn_sharedkv_metadata( + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_kv: int, + ori_topk: int, + cmp_topk: int, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + has_ori_kv: bool, + has_cmp_kv: bool, +) -> torch.Tensor: + """Build the AICPU tiling metadata for :func:`sparse_attn_sharedkv`.""" + return torch.ops.xllm_ops.sparse_attn_sharedkv_metadata( + num_heads_q, + num_heads_kv, + head_dim, + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + batch_size, + max_seqlen_q, + max_seqlen_kv, + ori_topk, + cmp_topk, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + has_ori_kv, + has_cmp_kv, + ) + + +def quant_lightning_indexer( + query: torch.Tensor, + key: torch.Tensor, + weights: torch.Tensor, + query_dequant_scale: torch.Tensor, + key_dequant_scale: torch.Tensor, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + block_table: torch.Tensor | None, + metadata: torch.Tensor | None, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + return_value: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select the compressed key blocks each query attends to (int8 q/k).""" + return torch.ops.xllm_ops.quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + query_quant_mode, + key_quant_mode, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + metadata, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + return_value, + ) + + +def quant_lightning_indexer_metadata( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_k: int, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + device: str, +) -> torch.Tensor: + """Build the AICPU tiling metadata for :func:`quant_lightning_indexer`.""" + return torch.ops.xllm_ops.quant_lightning_indexer_metadata( + num_heads_q, + num_heads_k, + head_dim, + query_quant_mode, + key_quant_mode, + actual_seq_lengths_query, + actual_seq_lengths_key, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + device, + ) diff --git a/xllm/python/kernels_npu/moe.py b/xllm/python/kernels_npu/moe.py index 9fdca9817b..5766ffaa97 100644 --- a/xllm/python/kernels_npu/moe.py +++ b/xllm/python/kernels_npu/moe.py @@ -107,23 +107,19 @@ def grouped_moe( num_tokens = hidden_states.shape[0] num_experts = gating_output.shape[1] expert_range = active_expert_range if active_expert_range is not None else [0, num_experts] - sorted_hidden_i8, expanded_row_idx, expert_tokens, pertoken_scale = ( - torch_npu.npu_moe_init_routing_v2( - hidden_states, - topk_ids.to(torch.int32), - scale=None, - active_num=num_tokens * topk, - expert_num=num_experts, - expert_tokens_num_type=1, - expert_tokens_num_flag=True, - active_expert_range=expert_range, - quant_mode=1, - ) + sorted_hidden_i8, expanded_row_idx, expert_tokens, pertoken_scale = torch_npu.npu_moe_init_routing_v2( + hidden_states, + topk_ids.to(torch.int32), + scale=None, + active_num=num_tokens * topk, + expert_num=num_experts, + expert_tokens_num_type=1, + expert_tokens_num_flag=True, + active_expert_range=expert_range, + quant_mode=1, ) num_local_experts = expert_range[1] - expert_range[0] - group_list = torch.cumsum( - expert_tokens[:num_local_experts].to(torch.int64), 0 - ) + group_list = torch.cumsum(expert_tokens[:num_local_experts].to(torch.int64), 0) act_i8, act_pt, _ = torch.ops.npu.npu_grouped_matmul_swiglu_quant( x=sorted_hidden_i8, weight=w13, @@ -152,6 +148,129 @@ def grouped_moe( ) +def _group_gemm(**kwargs) -> torch.Tensor: + return torch.ops.xllm_ops.group_gemm(**kwargs) + + +def _grouped_moe_with_selected_experts_impl( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + """Run grouped quantized experts with pre-computed routing (no gate). + + The routing and W8A8 grouped-matmul sequence mirrors the native NPU + ``FusedMoEImpl::select_experts`` and ``forward_expert`` paths. + """ + num_tokens = hidden_states.shape[0] + expert_num = num_total_experts if num_total_experts > 0 else w13.shape[0] + local_expert_count = num_experts_per_rank if num_experts_per_rank > 0 else expert_num + active_range = [start_expert_id, start_expert_id + local_expert_count] + expanded_hidden, expanded_row_idx, expert_tokens, _ = torch_npu.npu_moe_init_routing_v2( + hidden_states, + topk_ids.to(torch.int32), + scale=None, + active_num=num_tokens * topk_ids.size(-1), + expert_num=expert_num, + expert_tokens_num_type=1, + expert_tokens_num_flag=True, + active_expert_range=active_range, + quant_mode=-1, + ) + from xllm.python import kernels as _kernels + + sorted_hidden_i8, pertoken_scale = _kernels.dynamic_quant(expanded_hidden) + if pertoken_scale is None: + raise RuntimeError("dynamic_quant did not return a per-token scale") + group_list = expert_tokens.to(torch.int64) + gemm1_out = _group_gemm( + x=sorted_hidden_i8, + weight=w13, + scale=None, + per_token_scale=None, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=torch.int32, + ) + act_i8, act_pt = _kernels.dequant_swiglu_quant( + x=gemm1_out, + weight_scale=w13_scale, + activation_scale=pertoken_scale, + bias=None, + quant_scale=None, + quant_offset=None, + group_index=group_list.to(torch.int64), + activate_left=True, + quant_mode=1, + swiglu_mode=1, + clamp_limit=swiglu_limit, + glu_alpha=1.0, + glu_bias=0.0, + ) + del w13_offset, w2_offset + output = _group_gemm( + x=act_i8, + weight=w2, + scale=w2_scale.to(hidden_states.dtype), + per_token_scale=act_pt, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=hidden_states.dtype, + ) + return torch_npu.npu_moe_token_unpermute( + permuted_tokens=output, + sorted_indices=expanded_row_idx.abs(), + probs=topk_weights.to(output.dtype), + ) + + +@torch.library.custom_op("xllm_python::grouped_moe_with_selected_experts", mutates_args=()) +def grouped_moe_with_selected_experts( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + return _grouped_moe_with_selected_experts_impl( + hidden_states, + topk_weights, + topk_ids, + w13, + w2, + w13_scale, + w2_scale, + w13_offset, + w2_offset, + num_total_experts, + start_expert_id, + num_experts_per_rank, + swiglu_limit, + ) + + @grouped_moe.register_fake def _grouped_moe_fake( hidden_states: torch.Tensor, @@ -183,6 +302,27 @@ def _grouped_moe_fake( return torch.empty_like(hidden_states) +@grouped_moe_with_selected_experts.register_fake +def _grouped_moe_with_selected_experts_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + del topk_weights, topk_ids, w13, w2, w13_scale, w2_scale, w13_offset, w2_offset + del num_total_experts, start_expert_id, num_experts_per_rank, swiglu_limit + return torch.empty_like(hidden_states) + + def moe_fused_topk( gating_output: torch.Tensor, topk: int, @@ -202,8 +342,7 @@ def moe_fused_topk( """ del gating_output, topk, renormalize, scoring_func raise NotImplementedError( - "moe_fused_topk has no NPU kernel; NPU routes and runs experts in one " - "step through grouped_moe" + "moe_fused_topk has no NPU kernel; NPU routes and runs experts in one step through grouped_moe" ) @@ -245,10 +384,7 @@ def cutlass_fused_moe( ep_size, ep_rank, ) - raise NotImplementedError( - "cutlass_fused_moe is a CUDA library kernel; the NPU equivalent is " - "grouped_moe" - ) + raise NotImplementedError("cutlass_fused_moe is a CUDA library kernel; the NPU equivalent is grouped_moe") def fused_moe( @@ -272,8 +408,7 @@ def fused_moe( """ del hidden_states, topk_ids, topk_weights, w13, w2 raise NotImplementedError( - "fused_moe has no NPU kernel; see kernels_cuda/triton/fused_moe.py for " - "the reference implementation" + "fused_moe has no NPU kernel; see kernels_cuda/triton/fused_moe.py for the reference implementation" ) diff --git a/xllm/python/kernels_npu/normalization.py b/xllm/python/kernels_npu/normalization.py index e529abf6ef..899c97cb15 100644 --- a/xllm/python/kernels_npu/normalization.py +++ b/xllm/python/kernels_npu/normalization.py @@ -20,6 +20,7 @@ rms_norm = torch.ops.xllm_ops.rms_norm fused_add_rms_norm = torch.ops.xllm_ops.fused_add_rms_norm +rms_norm_dynamic_quant = torch.ops.xllm_ops.rms_norm_dynamic_quant def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: @@ -34,8 +35,7 @@ def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: """ del value, eps raise NotImplementedError( - "l2_norm has no NPU kernel; see kernels_cuda/triton/l2_norm.py for the " - "reference implementation" + "l2_norm has no NPU kernel; see kernels_cuda/triton/l2_norm.py for the reference implementation" ) @@ -58,9 +58,14 @@ def rms_norm_gated( """ 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" + "rms_norm_gated has no NPU kernel; see kernels_cuda/triton/rms_norm.py for the reference implementation" ) -__all__ = ["rms_norm", "fused_add_rms_norm", "l2_norm", "rms_norm_gated"] +__all__ = [ + "rms_norm", + "fused_add_rms_norm", + "rms_norm_dynamic_quant", + "l2_norm", + "rms_norm_gated", +] diff --git a/xllm/python/kernels_npu/rotary_embedding.py b/xllm/python/kernels_npu/rotary_embedding.py index 693176f663..4d793fcac1 100644 --- a/xllm/python/kernels_npu/rotary_embedding.py +++ b/xllm/python/kernels_npu/rotary_embedding.py @@ -76,9 +76,7 @@ def fused_qk_norm_rope( ) -@torch.library.custom_op( - "xllm_python::interleaved_rotary_embedding", mutates_args=() -) +@torch.library.custom_op("xllm_python::interleaved_rotary_embedding", mutates_args=()) def interleaved_rotary_embedding( value: torch.Tensor, cosine: torch.Tensor, @@ -95,9 +93,7 @@ def interleaved_rotary_embedding( A tensor with the shape and dtype of ``value``. """ num_tokens, num_heads, head_dim = value.shape - output = torch_npu.npu_interleave_rope( - value.view(num_tokens, num_heads, 1, head_dim), cosine, sine - ) + output = torch_npu.npu_interleave_rope(value.view(num_tokens, num_heads, 1, head_dim), cosine, sine) return output.view(num_tokens, num_heads, head_dim) @@ -170,9 +166,38 @@ def vision_rotary_mul( """ import torch_npu - return torch_npu.npu_rotary_mul( - value.unsqueeze(0).contiguous(), cos_full, sin_full - ).squeeze(0) + return torch_npu.npu_rotary_mul(value.unsqueeze(0).contiguous(), cos_full, sin_full).squeeze(0) + + +def npu_inplace_partial_rotary_mul( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_start_dim: int, + rope_head_dim: int, + inverse: bool = False, +) -> torch.Tensor: + """In-place partial interleaved RoPE on the ``[rope_start_dim:]`` slice. + + Mirrors C++ ``apply_partial_rope`` (deepseek_sparse_attention.cpp:151-190): + x is 3D ``[M, n_head, head_dim]``; cos/sin are 2D ``[M, rope_head_dim]`` + (per-token, no head dim). Reshaped to 4D for the NPU kernel + (``aclnnInplacePartialRotaryMul``, rotary_mode="interleave", + partial_slice=[rope_start_dim, rope_start_dim+rope_head_dim] -- a half-open + range, NOT [start, length]). Modifies x in place. + """ + x4d = x.unsqueeze(2) # [M, n_head, 1, head_dim] + cos4d = cos.view(cos.size(0), 1, 1, cos.size(1)) + sin_cache = -sin if inverse else sin + sin4d = sin_cache.view(sin.size(0), 1, 1, sin.size(1)) + torch.ops.xllm_ops.npu_inplace_partial_rotary_mul( + x4d, + cos4d, + sin4d, + "interleave", + [int(rope_start_dim), int(rope_start_dim + rope_head_dim)], + ) + return x __all__ = [ @@ -180,4 +205,5 @@ def vision_rotary_mul( "interleaved_rotary_embedding", "mrope", "vision_rotary_mul", + "npu_inplace_partial_rotary_mul", ] diff --git a/xllm/python/layers/embedding.py b/xllm/python/layers/embedding.py index b6db2fe4c9..e8f9559a61 100644 --- a/xllm/python/layers/embedding.py +++ b/xllm/python/layers/embedding.py @@ -52,5 +52,5 @@ def __init__( def forward(self, input_ids: torch.Tensor) -> torch.Tensor: out = torch.nn.functional.embedding(input_ids, self.weight) if self.tp_size > 1: - out = distributed.all_gather(out, dim=-1, world_size=self.tp_size) + out = distributed.tp_all_gather(out, dim=-1, world_size=self.tp_size) return out diff --git a/xllm/python/layers/linear.py b/xllm/python/layers/linear.py index e73f2f5d3f..b90be80837 100644 --- a/xllm/python/layers/linear.py +++ b/xllm/python/layers/linear.py @@ -62,16 +62,14 @@ def __init__( ) ) if bias: - self.bias = nn.Parameter( - torch.empty(out_features_per_partition, dtype=dtype, device=device) - ) + self.bias = nn.Parameter(torch.empty(out_features_per_partition, dtype=dtype, device=device)) else: self.register_parameter("bias", None) def forward(self, x: torch.Tensor) -> torch.Tensor: out = torch.nn.functional.linear(x, self.weight, self.bias) if self.gather_output and self.tp_size > 1: - out = distributed.all_gather(out, dim=-1, world_size=self.tp_size) + out = distributed.tp_all_gather(out, dim=-1, world_size=self.tp_size) return out @@ -101,9 +99,7 @@ def __init__( if bias and not reduce_results: # The bias is replicated and must be added exactly once, which is # only possible here when this layer owns the reduction. - raise ValueError( - "a deferred reduction cannot be combined with a replicated bias" - ) + raise ValueError("a deferred reduction cannot be combined with a replicated bias") self.weight = nn.Parameter( torch.empty( out_features, @@ -113,9 +109,7 @@ def __init__( ) ) if bias: - self.bias = nn.Parameter( - torch.empty(out_features, dtype=dtype, device=device) - ) + self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device)) else: self.register_parameter("bias", None) @@ -133,7 +127,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: else: out = torch.nn.functional.linear(x, self.weight) if self.tp_size > 1 and self.reduce_results: - distributed.all_reduce_(out) + distributed.tp_all_reduce(out) if self.bias is not None: out = out + self.bias return out diff --git a/xllm/python/model_executor/executor.py b/xllm/python/model_executor/executor.py index 656afa06a5..ad22d0a0ef 100644 --- a/xllm/python/model_executor/executor.py +++ b/xllm/python/model_executor/executor.py @@ -42,11 +42,31 @@ def _create_attention_backend( first_attention: Attention, device: torch.device, dtype: torch.dtype, + config: dict | None = None, ) -> AttentionBackend: + config = config or {} + model_type = config.get("model_type", "") + if model_type == "deepseek_v4" and current_platform.is_npu(): + from xllm.python.attention.dsa_attention import DsaAttentionBackend + + return DsaAttentionBackend( + compress_ratios=list(config.get("compress_ratios", [])), + window_size=int(config.get("window_size", 128)), + n_layers=int(config.get("n_layers", config.get("num_hidden_layers", 0))), + num_heads=first_attention.num_heads, + attn_head_dim=first_attention.head_dim, + index_topk=int(config.get("index_topk", 512)), + index_n_heads=int(config.get("index_n_heads", 64)), + index_head_dim=int(config.get("index_head_dim", 128)), + rope_head_dim=int(config.get("qk_rope_head_dim", 64)), + device=device, + dtype=dtype, + ) if current_platform.is_npu(): from xllm.python.attention.npu_paged_attention import ( NpuPagedAttentionBackend, ) + return NpuPagedAttentionBackend( num_heads=first_attention.num_heads, num_kv_heads=first_attention.num_kv_heads, @@ -58,6 +78,7 @@ def _create_attention_backend( ) if current_platform.is_cuda(): from xllm.python.attention.flashinfer import FlashInferBackend + return FlashInferBackend( num_heads=first_attention.num_heads, num_kv_heads=first_attention.num_kv_heads, @@ -67,9 +88,7 @@ def _create_attention_backend( device=device, dtype=dtype, ) - raise NotImplementedError( - f"No attention backend available for device type '{device.type}'" - ) + raise NotImplementedError(f"No attention backend available for device type '{device.type}'") class ModelExecutor: @@ -82,9 +101,7 @@ def __init__( self.model = model self._kv_bound = False - attention_layers = [ - module for module in model.modules() if isinstance(module, Attention) - ] + attention_layers = [module for module in model.modules() if isinstance(module, Attention)] if not attention_layers: raise ValueError("Python model does not contain an Attention layer") @@ -92,17 +109,12 @@ def __init__( expected_config = self._attention_config(first_attention) for layer in attention_layers[1:]: if self._attention_config(layer) != expected_config: - raise ValueError( - "Attention backend requires identical attention configuration " - "across all layers" - ) + raise ValueError("Attention backend requires identical attention configuration across all layers") first_parameter = next(model.parameters()) device = first_parameter.device self._num_attention_layers = len(attention_layers) - self.attention_backend = _create_attention_backend( - first_attention, device, first_parameter.dtype - ) + self.attention_backend = _create_attention_backend(first_attention, device, first_parameter.dtype, config) execution_model = model.model self.eager_runner = EagerRunner(execution_model, self.attention_backend, device) @@ -124,16 +136,14 @@ def __init__( "cudagraphs", "aclgraph", ): - raise NotImplementedError( - "Python data parallel graph execution supports cudagraphs and " - "aclgraph only" - ) + raise NotImplementedError("Python data parallel graph execution supports cudagraphs and aclgraph only") if graph_backend in ("", "off", "none", "0"): pass elif graph_backend == "cudagraphs": from xllm.python.model_executor.runners.decode_cuda_graph import ( DecodeCudaGraphRunner, ) + self.decode_graph_runner = DecodeCudaGraphRunner( execution_model, self.attention_backend, @@ -147,6 +157,7 @@ def __init__( from xllm.python.model_executor.runners.decode_acl_graph import ( DecodeAclGraphRunner, ) + self.decode_graph_runner = DecodeAclGraphRunner( execution_model, self.attention_backend, @@ -168,9 +179,8 @@ def __init__( "graph_backend=off/aclgraph, or set cp_size=1." ) from xllm.python.model_executor.runners.inductor import InductorRunner - self.inductor_runner = InductorRunner( - execution_model, self.attention_backend, device, graph_backend - ) + + self.inductor_runner = InductorRunner(execution_model, self.attention_backend, device, graph_backend) @staticmethod def _attention_config(layer: Attention) -> tuple[int, int, int, float, int]: @@ -184,15 +194,9 @@ def _attention_config(layer: Attention) -> tuple[int, int, int, float, int]: def bind_kv_caches(self, kv_caches: list[LayerCacheInput]) -> None: layer_caches = normalize_layer_caches(kv_caches) - required_layers = max( - layer.layer_id - for layer in self.model.modules() - if isinstance(layer, Attention) - ) + 1 + required_layers = max(layer.layer_id for layer in self.model.modules() if isinstance(layer, Attention)) + 1 if len(layer_caches) < required_layers: - raise ValueError( - "cache layer count does not match the model layer layout" - ) + raise ValueError("cache layer count does not match the model layer layout") if self._kv_bound: return self.attention_backend.bind_kv_caches(layer_caches) @@ -216,15 +220,9 @@ def execute( raise RuntimeError("KV caches are not bound") graph_runner = self.decode_graph_runner - if graph_runner is not None and graph_runner.can_execute( - input_ids, metadata, input_embedding - ): - graph_runner.warmup( - input_ids.device, input_ids.dtype, input_embedding - ) - return graph_runner.execute( - input_ids, positions, metadata, input_embedding - ) + if graph_runner is not None and graph_runner.can_execute(input_ids, metadata, input_embedding): + graph_runner.warmup(input_ids.device, input_ids.dtype, input_embedding) + return graph_runner.execute(input_ids, positions, metadata, input_embedding) if self.inductor_runner is not None: return self.inductor_runner.execute( input_ids, diff --git a/xllm/python/models/deepseek_v32.py b/xllm/python/models/deepseek_v32.py index 1cf6c05a79..2c2ec5b01a 100644 --- a/xllm/python/models/deepseek_v32.py +++ b/xllm/python/models/deepseek_v32.py @@ -67,9 +67,7 @@ def _yarn_find_correction_dim( base: float, max_position_embeddings: int, ) -> float: - return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( - 2 * math.log(base) - ) + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) def _yarn_find_correction_range( @@ -78,7 +76,7 @@ def _yarn_find_correction_range( dim: int, base: float, max_position_embeddings: int, -) -> Tuple[int, int]: +) -> tuple[int, int]: low = _yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) high = _yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) low = math.floor(low) @@ -86,9 +84,7 @@ def _yarn_find_correction_range( return max(low, 0), min(high, dim - 1) -def _yarn_linear_ramp_mask( - low: float, high: float, dim: int, dtype: torch.dtype, device: torch.device -) -> torch.Tensor: +def _yarn_linear_ramp_mask(low: float, high: float, dim: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: if low == high: high += 0.001 # Prevent singularity. linear = (torch.arange(dim, dtype=dtype, device=device) - low) / (high - low) @@ -97,7 +93,7 @@ def _yarn_linear_ramp_mask( def _gather_interleave_cos_sin( cos_sin_cache: torch.Tensor, positions: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Gather per-token cos/sin and double for ``npu_interleave_rope``.""" cos_sin = cos_sin_cache[positions] half = cos_sin.size(-1) // 2 @@ -107,16 +103,12 @@ def _gather_interleave_cos_sin( return cos, sin -def _interleave_rope_with( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> torch.Tensor: +def _interleave_rope_with(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: """Apply interleaved RoPE to ``[T, H, D]`` with precomputed cos/sin.""" return kernels.interleaved_rotary_embedding(x, cos, sin) -def _apply_half_rope( - cos_sin_cache: torch.Tensor, x: torch.Tensor, positions: torch.Tensor -) -> torch.Tensor: +def _apply_half_rope(cos_sin_cache: torch.Tensor, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: """Half-rotate RoPE (NeoX style) for ``[T, H, D]`` tensors.""" cos_sin = cos_sin_cache[positions] half = cos_sin.size(-1) // 2 @@ -160,17 +152,13 @@ def __init__( device=device, ) freqs = torch.outer(t, inv_freq) - rope_mscale = _yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale( - scaling_factor, mscale_all_dim - ) + rope_mscale = _yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale(scaling_factor, mscale_all_dim) cos = freqs.cos() * rope_mscale sin = freqs.sin() * rope_mscale cache = torch.cat([cos, sin], dim=-1) if dtype is not None: cache = cache.to(dtype) - self.register_buffer( - "cos_sin_cache", cache.contiguous(), persistent=False - ) + self.register_buffer("cos_sin_cache", cache.contiguous(), persistent=False) @staticmethod def _yarn_inv_freq( @@ -182,10 +170,7 @@ def _yarn_inv_freq( max_position_embeddings: int, device: torch.device, ) -> torch.Tensor: - pos_freqs = base ** ( - torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) - / rotary_dim - ) + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) / rotary_dim) inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs) low, high = _yarn_find_correction_range( @@ -195,12 +180,7 @@ def _yarn_inv_freq( base, max_position_embeddings, ) - inv_freq_mask = ( - 1 - - _yarn_linear_ramp_mask( - low, high, rotary_dim // 2, torch.float32, device - ) - ) + inv_freq_mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float32, device) return inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask @@ -254,7 +234,7 @@ class DeepseekV3Config: world_size: int = 1 @classmethod - def from_dict(cls, d: dict) -> "DeepseekV3Config": + def from_dict(cls, d: dict) -> DeepseekV3Config: def pick(*keys, default=None): for k in keys: if k in d and d[k] is not None: @@ -286,12 +266,8 @@ def rpick(*keys, default=None): vocab_size=int(pick("vocab_size", default=129280)), rms_norm_eps=float(pick("rms_norm_eps", default=1e-6)), rope_theta=float(pick("rope_theta", default=1.0e6)), - max_position_embeddings=int( - pick("max_position_embeddings", default=4096) - ), - original_max_position_embeddings=int( - rpick("original_max_position_embeddings", default=4096) - ), + max_position_embeddings=int(pick("max_position_embeddings", default=4096)), + original_max_position_embeddings=int(rpick("original_max_position_embeddings", default=4096)), rope_scaling_factor=float(rpick("factor", "rope_scaling_factor", default=40.0)), rope_beta_fast=int(rpick("beta_fast", default=32)), rope_beta_slow=int(rpick("beta_slow", default=1)), @@ -316,9 +292,7 @@ def rpick(*keys, default=None): routed_scaling_factor=float(pick("routed_scaling_factor", default=2.5)), topk_method=str(pick("topk_method", default="noaux_tc")), norm_topk_prob=bool(pick("norm_topk_prob", default=True)), - moe_intermediate_size=int( - pick("moe_intermediate_size", default=2048) - ), + moe_intermediate_size=int(pick("moe_intermediate_size", default=2048)), tp_size=int(pick("tp_size", default=1)), tp_rank=int(pick("tp_rank", default=0)), ep_size=int(pick("ep_size", default=1)), @@ -330,33 +304,28 @@ def rpick(*keys, default=None): world_size=int(pick("world_size", default=1)), ) - def head_split(self) -> Tuple[int, int]: + def head_split(self) -> tuple[int, int]: """Per-rank (num_heads_local, num_kv_heads_local=1).""" num_heads_local = self.n_heads // self.tp_size return num_heads_local, 1 def validate(self) -> None: if self.ep_size not in (1, self.world_size): - raise ValueError( - f"ep_size must be 1 or world_size ({self.world_size}), got {self.ep_size}" - ) + raise ValueError(f"ep_size must be 1 or world_size ({self.world_size}), got {self.ep_size}") if self.ep_size > 1 and self.n_routed_experts % self.ep_size: raise ValueError( - f"n_routed_experts ({self.n_routed_experts}) must be divisible by " - f"ep_size ({self.ep_size})" + f"n_routed_experts ({self.n_routed_experts}) must be divisible by ep_size ({self.ep_size})" ) if self.ep_size > 1 and self.moe_tp_size * self.ep_size != self.world_size: raise ValueError( - f"world_size ({self.world_size}) must equal " - f"moe_tp_size ({self.moe_tp_size}) * ep_size ({self.ep_size})" + f"world_size ({self.world_size}) must equal moe_tp_size ({self.moe_tp_size}) * ep_size ({self.ep_size})" ) class W8A8StaticLinear(nn.Module): """Static-activation W8A8 linear (attention projections).""" - def __init__(self, in_features: int, out_features: int, device: torch.device, - row_parallel: bool = False) -> None: + def __init__(self, in_features: int, out_features: int, device: torch.device, row_parallel: bool = False) -> None: super().__init__() self.in_features = in_features self.out_features = out_features @@ -365,18 +334,10 @@ def __init__(self, in_features: int, out_features: int, device: torch.device, torch.empty(out_features, in_features, dtype=torch.int8, device=device), requires_grad=False, ) - self.register_buffer( - "deq_scale", torch.empty(out_features, dtype=torch.float32, device=device) - ) - self.register_buffer( - "quant_bias", torch.empty(out_features, dtype=torch.int32, device=device) - ) - self.register_buffer( - "input_scale", torch.empty(1, dtype=torch.bfloat16, device=device) - ) - self.register_buffer( - "input_offset", torch.empty(1, dtype=torch.bfloat16, device=device) - ) + self.register_buffer("deq_scale", torch.empty(out_features, dtype=torch.float32, device=device)) + self.register_buffer("quant_bias", torch.empty(out_features, dtype=torch.int32, device=device)) + self.register_buffer("input_scale", torch.empty(1, dtype=torch.bfloat16, device=device)) + self.register_buffer("input_offset", torch.empty(1, dtype=torch.bfloat16, device=device)) def process_weights_after_loading(self) -> None: self.weight.data = self.weight.data.transpose(0, 1).contiguous() @@ -385,13 +346,17 @@ def process_weights_after_loading(self) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: mult = self.input_scale_recip x_int8 = torch.clamp( - torch.round( - x.to(torch.float32) * mult + self.input_offset.to(torch.float32) - ), - -128, 127, + torch.round(x.to(torch.float32) * mult + self.input_offset.to(torch.float32)), + -128, + 127, ).to(torch.int8) return kernels.quant_matmul( - x_int8, self.weight, False, self.deq_scale, None, None, + x_int8, + self.weight, + False, + self.deq_scale, + None, + None, self.quant_bias if not (self.row_parallel and distributed.tp_rank(x.device) != 0) else None, torch.bfloat16, ) @@ -400,27 +365,34 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class W8A8DynamicLinear(nn.Module): """Dynamic-activation W8A8 linear (MLP / experts).""" - def __init__(self, in_features: int, out_features: int, device: torch.device) -> None: + def __init__( + self, + in_features: int, + out_features: int, + device: torch.device, + transpose_weight_after_loading: bool = True, + ) -> None: super().__init__() self.in_features = in_features self.out_features = out_features + self.transpose_weight_after_loading = transpose_weight_after_loading + self._weight_is_transposed = False self.weight = nn.Parameter( torch.empty(out_features, in_features, dtype=torch.int8, device=device), requires_grad=False, ) - self.register_buffer( - "weight_scale", torch.empty(out_features, 1, dtype=torch.float32, device=device) - ) - self.register_buffer( - "weight_offset", torch.empty(out_features, 1, dtype=torch.float32, device=device) - ) + self.register_buffer("weight_scale", torch.empty(out_features, 1, dtype=torch.float32, device=device)) + self.register_buffer("weight_offset", torch.empty(out_features, 1, dtype=torch.float32, device=device)) def process_weights_after_loading(self) -> None: - self.weight.data = self.weight.data.transpose(0, 1).contiguous() + if self.transpose_weight_after_loading and not self._weight_is_transposed: + self.weight.data = self.weight.data.transpose(0, 1).contiguous() + self._weight_is_transposed = True self.weight_scale.data = self.weight_scale.data.flatten().contiguous() self.weight_offset.data = self.weight_offset.data.flatten().contiguous() if not bool(torch.all(self.weight_offset == 0)): import logging + logging.getLogger(__name__).warning( "W8A8DynamicLinear loaded with non-zero weight_offset; the " "int8 matmul path drops the antiquant offset -- output may be " @@ -430,11 +402,37 @@ def process_weights_after_loading(self) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: x_int8, pertoken = kernels.dynamic_quant(x) return kernels.quant_matmul( - x_int8, self.weight, False, self.weight_scale, None, - pertoken, None, torch.bfloat16, + x_int8, + self.weight, + not self._weight_is_transposed, + self.weight_scale, + None, + pertoken, + None, + torch.bfloat16, + ) + + def forward_quant(self, x_int8: torch.Tensor, pertoken: torch.Tensor) -> torch.Tensor: + """Run quantized matmul on an already-quantized activation.""" + return kernels.quant_matmul( + x_int8, + self.weight, + not self._weight_is_transposed, + self.weight_scale, + None, + pertoken, + None, + torch.bfloat16, ) +def _swiglu_with_clamp(x: torch.Tensor, limit: float) -> torch.Tensor: + gate, up = x.chunk(2, dim=-1) + gate = gate.to(torch.float32).clamp_max(limit) + up = up.to(torch.float32).clamp(min=-limit, max=limit) + return (torch.nn.functional.silu(gate) * up).to(x.dtype) + + class W8A8WeightLoader: """Shared W8A8 weight-loading helpers for a model's ``load_weights``. @@ -447,7 +445,7 @@ class W8A8WeightLoader: def __init__( self, model: nn.Module, - state_dicts: List["StateDict"], + state_dicts: list[StateDict], tp_size: int, tp_rank: int, ) -> None: @@ -457,7 +455,7 @@ def __init__( self.tp_size = tp_size self.tp_rank = tp_rank - def find(self, name: str) -> Optional["StateDict"]: + def find(self, name: str) -> Optional[StateDict]: for sd in self._state_dicts: if sd.has(name): return sd @@ -469,7 +467,10 @@ def load_tensor(self, name: str) -> torch.Tensor: return sd.get_tensor(name) def shard( - self, t: torch.Tensor, dim: int, world: Optional[int] = None, + self, + t: torch.Tensor, + dim: int, + world: Optional[int] = None, rank: Optional[int] = None, ) -> torch.Tensor: world = self.tp_size if world is None else world @@ -486,11 +487,8 @@ def copy_in(self, param_name: str, tensor: torch.Tensor) -> None: assert p is not None, f"no parameter/buffer named {param_name}" p.data.copy_(tensor.to(dtype=p.dtype, device=p.device)) - def load_w8a8_a( - self, prefix: str, proj: str, shard_dims: Optional[dict] = None - ) -> None: - for suffix in ("weight", "deq_scale", "quant_bias", - "input_scale", "input_offset"): + def load_w8a8_a(self, prefix: str, proj: str, shard_dims: Optional[dict] = None) -> None: + for suffix in ("weight", "deq_scale", "quant_bias", "input_scale", "input_offset"): t = self.load_tensor(prefix + proj + "." + suffix) dim = (shard_dims or {}).get(suffix) if dim is not None: @@ -504,18 +502,19 @@ def load_w8a8_b(self, mlp_pfx: str) -> None: uw = self.load_tensor(mlp_pfx + "up_proj.weight") us = self.load_tensor(mlp_pfx + "up_proj.weight_scale") uo = self.load_tensor(mlp_pfx + "up_proj.weight_offset") - self.copy_in(mlp_pfx + "gate_up_proj.weight", - torch.cat([self.shard(gw, 0), self.shard(uw, 0)], dim=0).contiguous()) - self.copy_in(mlp_pfx + "gate_up_proj.weight_scale", - torch.cat([self.shard(gs, 0), self.shard(us, 0)], dim=0).contiguous()) - self.copy_in(mlp_pfx + "gate_up_proj.weight_offset", - torch.cat([self.shard(go, 0), self.shard(uo, 0)], dim=0).contiguous()) - self.copy_in(mlp_pfx + "down_proj.weight", - self.shard(self.load_tensor(mlp_pfx + "down_proj.weight"), dim=1)) - self.copy_in(mlp_pfx + "down_proj.weight_scale", - self.load_tensor(mlp_pfx + "down_proj.weight_scale")) - self.copy_in(mlp_pfx + "down_proj.weight_offset", - self.load_tensor(mlp_pfx + "down_proj.weight_offset")) + self.copy_in( + mlp_pfx + "gate_up_proj.weight", torch.cat([self.shard(gw, 0), self.shard(uw, 0)], dim=0).contiguous() + ) + self.copy_in( + mlp_pfx + "gate_up_proj.weight_scale", torch.cat([self.shard(gs, 0), self.shard(us, 0)], dim=0).contiguous() + ) + self.copy_in( + mlp_pfx + "gate_up_proj.weight_offset", + torch.cat([self.shard(go, 0), self.shard(uo, 0)], dim=0).contiguous(), + ) + self.copy_in(mlp_pfx + "down_proj.weight", self.shard(self.load_tensor(mlp_pfx + "down_proj.weight"), dim=1)) + self.copy_in(mlp_pfx + "down_proj.weight_scale", self.load_tensor(mlp_pfx + "down_proj.weight_scale")) + self.copy_in(mlp_pfx + "down_proj.weight_offset", self.load_tensor(mlp_pfx + "down_proj.weight_offset")) class DeepseekV3MLP(nn.Module): @@ -529,15 +528,15 @@ def __init__( device: torch.device, skip_tp_reduce: bool = False, tp_override: Optional[int] = None, + swiglu_limit: float = 0.0, ) -> None: super().__init__() tp = tp_override if tp_override is not None else cfg.tp_size - assert intermediate_size % tp == 0, ( - f"intermediate_size {intermediate_size} not divisible by tp {tp}" - ) + assert intermediate_size % tp == 0, f"intermediate_size {intermediate_size} not divisible by tp {tp}" inter_local = intermediate_size // tp self.tp = tp self.skip_tp_reduce = skip_tp_reduce + self.swiglu_limit = swiglu_limit self.gate_up_proj = W8A8DynamicLinear(cfg.hidden_size, 2 * inter_local, device) self.down_proj = W8A8DynamicLinear(inter_local, cfg.hidden_size, device) @@ -547,7 +546,11 @@ def process_weights_after_loading(self) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up = self.gate_up_proj(x) - act = kernels.silu_and_mul(gate_up) + act = ( + _swiglu_with_clamp(gate_up, self.swiglu_limit) + if 0.0 < self.swiglu_limit < 1_000_000.0 + else kernels.silu_and_mul(gate_up) + ) out = self.down_proj(act) if self.tp > 1 and not self.skip_tp_reduce: distributed.all_reduce_(out) @@ -572,9 +575,7 @@ def __init__( qk_rope = cfg.qk_rope_head_dim v_head = cfg.v_head_dim scale = (qk_nope + qk_rope) ** -0.5 - attn_mscale = _yarn_get_mscale( - cfg.rope_scaling_factor, cfg.rope_mscale_all_dim - ) + attn_mscale = _yarn_get_mscale(cfg.rope_scaling_factor, cfg.rope_mscale_all_dim) scale = scale * attn_mscale * attn_mscale super().__init__( num_heads=num_heads, @@ -593,15 +594,9 @@ def __init__( self.q_a_proj = W8A8StaticLinear(cfg.hidden_size, cfg.q_lora_rank, device) self.kv_a_proj_with_mqa = W8A8StaticLinear(cfg.hidden_size, kv_lora + qk_rope, device) - self.q_a_layernorm = RMSNorm( - cfg.q_lora_rank, cfg.rms_norm_eps, dtype=dtype, device=device - ) - self.kv_a_layernorm = RMSNorm( - kv_lora, cfg.rms_norm_eps, dtype=dtype, device=device - ) - self.q_b_proj = W8A8StaticLinear( - cfg.q_lora_rank, num_heads * (qk_nope + qk_rope), device - ) + self.q_a_layernorm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps, dtype=dtype, device=device) + self.kv_a_layernorm = RMSNorm(kv_lora, cfg.rms_norm_eps, dtype=dtype, device=device) + self.q_b_proj = W8A8StaticLinear(cfg.q_lora_rank, num_heads * (qk_nope + qk_rope), device) self.kv_b_proj = ColumnParallelLinear( kv_lora, num_heads * (qk_nope + v_head), @@ -609,8 +604,7 @@ def __init__( dtype=dtype, device=device, ) - self.o_proj = W8A8StaticLinear(num_heads * v_head, cfg.hidden_size, device, - row_parallel=True) + self.o_proj = W8A8StaticLinear(num_heads * v_head, cfg.hidden_size, device, row_parallel=True) self.register_buffer( "W_UK", torch.empty(num_heads, qk_nope, kv_lora, dtype=dtype, device=device), @@ -621,9 +615,7 @@ def __init__( torch.empty(num_heads, kv_lora, v_head, dtype=dtype, device=device), persistent=False, ) - self.indexer: DeepseekV3Indexer | None = ( - DeepseekV3Indexer(cfg, dtype, device) if cfg.index_topk > 0 else None - ) + self.indexer: DeepseekV3Indexer | None = DeepseekV3Indexer(cfg, dtype, device) if cfg.index_topk > 0 else None def process_weights_after_loading(self) -> None: self.q_a_proj.process_weights_after_loading() @@ -636,9 +628,7 @@ def process_weights_after_loading(self) -> None: self.qk_nope_head_dim + self.v_head_dim, self.kv_lora_rank, ) - w_uk, w_uv = w.split( - [self.qk_nope_head_dim, self.v_head_dim], dim=1 - ) + w_uk, w_uv = w.split([self.qk_nope_head_dim, self.v_head_dim], dim=1) self.W_UK.copy_(w_uk.contiguous()) self.W_UV.copy_(w_uv.transpose(1, 2).contiguous()) @@ -655,41 +645,27 @@ def forward( topk = None if self.indexer is not None: ctx = backend.mla_index_context(self) - topk = self.indexer.select_qli( - hidden, q_c, positions, ctx, cos_sin_cache - ) + topk = self.indexer.select_qli(hidden, q_c, positions, ctx, cos_sin_cache) q = self.q_b_proj(q_c) q = q.view( num_tokens, self.num_heads_local, self.qk_nope_head_dim + self.qk_rope_head_dim, ) - q_nope, q_rope = q.split( - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) - q_latent = torch.bmm( - q_nope.transpose(0, 1), self.W_UK - ).transpose(0, 1) + q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_latent = torch.bmm(q_nope.transpose(0, 1), self.W_UK).transpose(0, 1) cos, sin = _gather_interleave_cos_sin(cos_sin_cache, positions) q_pe = _interleave_rope_with(q_rope, cos, sin) kv = self.kv_a_proj_with_mqa(hidden) - k_latent_raw, k_rope_raw = kv.split( - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) + k_latent_raw, k_rope_raw = kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_latent = self.kv_a_layernorm(k_latent_raw) k_pe = _interleave_rope_with(k_rope_raw.unsqueeze(1), cos, sin) k_latent_3d = k_latent.view(num_tokens, 1, self.kv_lora_rank) k_pe_3d = k_pe.view(num_tokens, 1, self.qk_rope_head_dim) - attn_out = backend.execute_mla( - q_latent, q_pe, k_latent_3d, k_pe_3d, self, topk=topk - ) - v_full = torch.bmm( - attn_out.transpose(0, 1), self.W_UV - ).transpose(0, 1) - v_full = v_full.reshape( - num_tokens, self.num_heads_local * self.v_head_dim - ) + attn_out = backend.execute_mla(q_latent, q_pe, k_latent_3d, k_pe_3d, self, topk=topk) + v_full = torch.bmm(attn_out.transpose(0, 1), self.W_UV).transpose(0, 1) + v_full = v_full.reshape(num_tokens, self.num_heads_local * self.v_head_dim) o = self.o_proj(v_full) if self.cfg.tp_size > 1: distributed.all_reduce_(o) @@ -699,21 +675,16 @@ def forward( class DeepseekV3Indexer(nn.Module): """DeepSeek-V3.2 lightning indexer (bf16 weights, non-quant aclnnLightningIndexer).""" - def __init__(self, cfg: "DeepseekV3Config", dtype: torch.dtype, - device: torch.device) -> None: + def __init__(self, cfg: DeepseekV3Config, dtype: torch.dtype, device: torch.device) -> None: super().__init__() self.n_head = cfg.index_n_heads self.head_dim = cfg.index_head_dim self.rope_dim = cfg.qk_rope_head_dim self.topk = cfg.index_topk - self.wq_b = nn.Linear(cfg.q_lora_rank, self.n_head * self.head_dim, - bias=False, dtype=dtype, device=device) - self.wk = nn.Linear(cfg.hidden_size, self.head_dim, - bias=False, dtype=dtype, device=device) - self.weights_proj = nn.Linear(cfg.hidden_size, self.n_head, - bias=False, dtype=dtype, device=device) - self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6, - dtype=dtype, device=device) + self.wq_b = nn.Linear(cfg.q_lora_rank, self.n_head * self.head_dim, bias=False, dtype=dtype, device=device) + self.wk = nn.Linear(cfg.hidden_size, self.head_dim, bias=False, dtype=dtype, device=device) + self.weights_proj = nn.Linear(cfg.hidden_size, self.n_head, bias=False, dtype=dtype, device=device) + self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6, dtype=dtype, device=device) def select_qli( self, @@ -724,19 +695,13 @@ def select_qli( cos_sin_cache: torch.Tensor, ) -> torch.Tensor: q = self.wq_b(qr).view(-1, self.n_head, self.head_dim) - q_pe, q_nope = torch.split( - q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) + q_pe, q_nope = torch.split(q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1) k = self.wk(hidden) weights = self.weights_proj(hidden) k = self.k_norm(k) - k_pe, k_nope = torch.split( - k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) + k_pe, k_nope = torch.split(k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1) q_pe = _apply_half_rope(cos_sin_cache, q_pe, positions) - k_pe = _apply_half_rope( - cos_sin_cache, k_pe.unsqueeze(1), positions - ).squeeze(1) + k_pe = _apply_half_rope(cos_sin_cache, k_pe.unsqueeze(1), positions).squeeze(1) q = torch.cat([q_pe, q_nope], dim=-1) k = torch.cat([k_pe, k_nope], dim=-1) if ctx.index_cache is not None and ctx.slot_mapping is not None: @@ -747,22 +712,28 @@ def select_qli( buffer_key = tuple(output_shape) topk_buffer = get_execution_buffer( ("LIGHTNING_INDEXER_INDICES",) + buffer_key, - lambda: torch.empty( - output_shape, dtype=torch.int32, device=q.device - ), + lambda: torch.empty(output_shape, dtype=torch.int32, device=q.device), ) values_buffer = get_execution_buffer( ("LIGHTNING_INDEXER_VALUES",) + buffer_key, - lambda: torch.empty( - output_shape, dtype=q.dtype, device=q.device - ), + lambda: torch.empty(output_shape, dtype=q.dtype, device=q.device), ) topk = kernels.lightning_indexer_out( - q, ctx.index_cache, weights, - ctx.actual_seq_q, ctx.actual_seq_kv, ctx.block_table, - "TND", "PA_BSND", self.topk, 3, - 9223372036854775807, 9223372036854775807, - False, topk_buffer, values_buffer, + q, + ctx.index_cache, + weights, + ctx.actual_seq_q, + ctx.actual_seq_kv, + ctx.block_table, + "TND", + "PA_BSND", + self.topk, + 3, + 9223372036854775807, + 9223372036854775807, + False, + topk_buffer, + values_buffer, ) return topk @@ -802,9 +773,7 @@ def __init__( self.local_expert_start = self.ep_rank * num_local_experts self.local_expert_end = self.local_expert_start + num_local_experts - self.gate = nn.Linear( - cfg.hidden_size, self.num_experts, bias=False, dtype=dtype, device=device - ) + self.gate = nn.Linear(cfg.hidden_size, self.num_experts, bias=False, dtype=dtype, device=device) self.register_buffer( "e_score_correction_bias", torch.zeros(self.num_experts, dtype=torch.float32, device=device), @@ -812,79 +781,87 @@ def __init__( ) self.experts_w13 = nn.Parameter( torch.empty( - num_local_experts, 2 * self.inter_local, self.hidden, - dtype=torch.int8, device=device, + num_local_experts, + 2 * self.inter_local, + self.hidden, + dtype=torch.int8, + device=device, ), requires_grad=False, ) self.register_buffer( "experts_w13_scale", torch.empty( - num_local_experts, 2 * self.inter_local, 1, - dtype=torch.float32, device=device, + num_local_experts, + 2 * self.inter_local, + 1, + dtype=torch.float32, + device=device, ), ) self.register_buffer( "experts_w13_offset", torch.empty( - num_local_experts, 2 * self.inter_local, 1, - dtype=torch.float32, device=device, + num_local_experts, + 2 * self.inter_local, + 1, + dtype=torch.float32, + device=device, ), ) self.experts_w2 = nn.Parameter( torch.empty( - num_local_experts, self.hidden, self.inter_local, - dtype=torch.int8, device=device, + num_local_experts, + self.hidden, + self.inter_local, + dtype=torch.int8, + device=device, ), requires_grad=False, ) self.register_buffer( "experts_w2_scale", torch.empty( - num_local_experts, self.hidden, 1, - dtype=torch.float32, device=device, + num_local_experts, + self.hidden, + 1, + dtype=torch.float32, + device=device, ), ) self.register_buffer( "experts_w2_offset", torch.empty( - num_local_experts, self.hidden, 1, - dtype=torch.float32, device=device, + num_local_experts, + self.hidden, + 1, + dtype=torch.float32, + device=device, ), ) shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts shared_tp = cfg.moe_tp_size if cfg.ep_size > 1 else None - self.shared_experts = DeepseekV3MLP(cfg, shared_inter, dtype, device, - skip_tp_reduce=True, - tp_override=shared_tp) + self.shared_experts = DeepseekV3MLP( + cfg, shared_inter, dtype, device, skip_tp_reduce=True, tp_override=shared_tp + ) def process_weights_after_loading(self) -> None: assert torch.all(self.experts_w13_offset == 0), ( - "DeepseekV3MoE int8-grouped path needs symmetric int8 experts " - "(experts_w13_offset == 0)") + "DeepseekV3MoE int8-grouped path needs symmetric int8 experts (experts_w13_offset == 0)" + ) assert torch.all(self.experts_w2_offset == 0), ( - "DeepseekV3MoE int8-grouped path needs symmetric int8 experts " - "(experts_w2_offset == 0)") + "DeepseekV3MoE int8-grouped path needs symmetric int8 experts (experts_w2_offset == 0)" + ) self.experts_w13.data = self.experts_w13.data.transpose(1, 2).contiguous() self.experts_w2.data = self.experts_w2.data.transpose(1, 2).contiguous() - self.experts_w13.data, self.experts_w2.data = ( - kernels.prepare_grouped_moe_weights( - self.experts_w13.data, - self.experts_w2.data, - ) + self.experts_w13.data, self.experts_w2.data = kernels.prepare_grouped_moe_weights( + self.experts_w13.data, + self.experts_w2.data, ) - self.experts_w13_scale.data = self.experts_w13_scale.data.view( - self.num_local_experts, -1 - ).contiguous() - self.experts_w13_offset.data = self.experts_w13_offset.data.view( - self.num_local_experts, -1 - ).contiguous() - self.experts_w2_scale.data = self.experts_w2_scale.data.view( - self.num_local_experts, -1 - ).contiguous() - self.experts_w2_offset.data = self.experts_w2_offset.data.view( - self.num_local_experts, -1 - ).contiguous() + self.experts_w13_scale.data = self.experts_w13_scale.data.view(self.num_local_experts, -1).contiguous() + self.experts_w13_offset.data = self.experts_w13_offset.data.view(self.num_local_experts, -1).contiguous() + self.experts_w2_scale.data = self.experts_w2_scale.data.view(self.num_local_experts, -1).contiguous() + self.experts_w2_offset.data = self.experts_w2_offset.data.view(self.num_local_experts, -1).contiguous() self.shared_experts.process_weights_after_loading() def forward(self, hidden: torch.Tensor) -> torch.Tensor: @@ -893,7 +870,10 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: if self.dp_size > 1: token_counts = list(get_forward_context().metadata.dp_token_counts) hidden = distributed.all_gather_variable( - hidden, token_counts, self.dp_rank, "dp", + hidden, + token_counts, + self.dp_rank, + "dp", ) logits = self.gate(hidden) @@ -927,7 +907,7 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: local_tokens = token_counts[self.dp_rank] if local_tokens == 0: return torch.zeros_like(local_hidden) - start = sum(token_counts[:self.dp_rank]) + start = sum(token_counts[: self.dp_rank]) final = final.narrow(0, start, local_tokens) return final @@ -943,17 +923,11 @@ def __init__( ) -> None: super().__init__() self.layer_id = layer_id - self.input_layernorm = RMSNorm( - cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device - ) + self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) self.self_attn = DeepseekV3MLAAttention(cfg, layer_id, dtype, device) - self.post_attention_layernorm = RMSNorm( - cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device - ) + self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) if layer_id < cfg.first_k_dense_replace: - self.mlp = DeepseekV3MLP( - cfg, cfg.intermediate_size, dtype, device - ) + self.mlp = DeepseekV3MLP(cfg, cfg.intermediate_size, dtype, device) else: self.mlp = DeepseekV3MoE(cfg, layer_id, dtype, device) @@ -963,7 +937,7 @@ def forward( residual: Optional[torch.Tensor], positions: torch.Tensor, cos_sin_cache: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: if residual is None: residual = hidden hidden = self.input_layernorm(hidden) @@ -976,9 +950,7 @@ def forward( class DeepseekV3Model(nn.Module): - def __init__( - self, cfg: DeepseekV3Config, dtype: torch.dtype, device: torch.device - ) -> None: + def __init__(self, cfg: DeepseekV3Config, dtype: torch.dtype, device: torch.device) -> None: super().__init__() tp = cfg.tp_size assert cfg.hidden_size % tp == 0 @@ -990,15 +962,8 @@ def __init__( dtype=dtype, device=device, ) - self.layers = nn.ModuleList( - [ - DeepseekV3DecoderLayer(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 - ) + self.layers = nn.ModuleList([DeepseekV3DecoderLayer(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) self.rotary = DeepseekYarnRotaryEmbedding( cfg.qk_rope_head_dim, cfg.original_max_position_embeddings, @@ -1012,9 +977,7 @@ def __init__( device=device, ) - def forward( - self, input_ids: torch.Tensor, positions: torch.Tensor - ) -> torch.Tensor: + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: hidden = self.embed_tokens(input_ids) positions = positions.to(torch.int64).contiguous() cos_sin_cache = self.rotary.cos_sin_cache @@ -1032,8 +995,7 @@ def __init__(self, config: dict, build_model: bool = True) -> None: super().__init__() self.cfg = DeepseekV3Config.from_dict(config) self.cfg.tp_size = int(config.get("tp_size", 1)) - self.cfg.tp_rank = int(config.get( - "tp_rank", _tp_rank_from_device(config.get("device", "npu:0")))) + self.cfg.tp_rank = int(config.get("tp_rank", _tp_rank_from_device(config.get("device", "npu:0")))) self.cfg.ep_size = int(config.get("ep_size", 1)) self.cfg.ep_rank = int(config.get("ep_rank", 0)) self.cfg.dp_size = int(config.get("dp_size", 1)) @@ -1043,9 +1005,7 @@ def __init__(self, config: dict, build_model: bool = True) -> None: self.cfg.world_size = int(config.get("world_size", self.cfg.tp_size)) if hasattr(self.cfg, "validate"): self.cfg.validate() - dtype = self.resolve_dtype( - config.get("dtype") or config.get("torch_dtype") - ) + dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype")) device = torch.device(config.get("device", "cuda")) self.dtype = dtype self.device = device @@ -1082,40 +1042,32 @@ def load_weights( if load_embedding: loader.copy_in( "model.embed_tokens.weight", - loader.shard( - loader.load_tensor("model.embed_tokens.weight"), dim=1 - ), + loader.shard(loader.load_tensor("model.embed_tokens.weight"), dim=1), ) for i in range(cfg.n_layers): p = f"model.layers.{i}." - loader.copy_in(p + "input_layernorm.weight", - loader.load_tensor(p + "input_layernorm.weight")) - loader.copy_in(p + "post_attention_layernorm.weight", - loader.load_tensor(p + "post_attention_layernorm.weight")) + loader.copy_in(p + "input_layernorm.weight", loader.load_tensor(p + "input_layernorm.weight")) + loader.copy_in( + p + "post_attention_layernorm.weight", loader.load_tensor(p + "post_attention_layernorm.weight") + ) attn = p + "self_attn." loader.load_w8a8_a(attn, "q_a_proj") - loader.copy_in(attn + "q_a_layernorm.weight", - loader.load_tensor(attn + "q_a_layernorm.weight")) - loader.load_w8a8_a(attn, "q_b_proj", - {"weight": 0, "deq_scale": 0, "quant_bias": 0}) + loader.copy_in(attn + "q_a_layernorm.weight", loader.load_tensor(attn + "q_a_layernorm.weight")) + loader.load_w8a8_a(attn, "q_b_proj", {"weight": 0, "deq_scale": 0, "quant_bias": 0}) loader.load_w8a8_a(attn, "kv_a_proj_with_mqa") - loader.copy_in(attn + "kv_a_layernorm.weight", - loader.load_tensor(attn + "kv_a_layernorm.weight")) - loader.copy_in(attn + "kv_b_proj.weight", - loader.shard(loader.load_tensor(attn + "kv_b_proj.weight"), dim=0)) + loader.copy_in(attn + "kv_a_layernorm.weight", loader.load_tensor(attn + "kv_a_layernorm.weight")) + loader.copy_in( + attn + "kv_b_proj.weight", loader.shard(loader.load_tensor(attn + "kv_b_proj.weight"), dim=0) + ) loader.load_w8a8_a(attn, "o_proj", {"weight": 1}) if cfg.index_topk > 0: idx = attn + "indexer." - loader.copy_in(idx + "wq_b.weight", - loader.load_tensor(idx + "wq_b.weight")) + loader.copy_in(idx + "wq_b.weight", loader.load_tensor(idx + "wq_b.weight")) loader.copy_in(idx + "wk.weight", loader.load_tensor(idx + "wk.weight")) - loader.copy_in(idx + "weights_proj.weight", - loader.load_tensor(idx + "weights_proj.weight")) - loader.copy_in(idx + "k_norm.weight", - loader.load_tensor(idx + "k_norm.weight")) - loader.copy_in(idx + "k_norm.bias", - loader.load_tensor(idx + "k_norm.bias")) + loader.copy_in(idx + "weights_proj.weight", loader.load_tensor(idx + "weights_proj.weight")) + loader.copy_in(idx + "k_norm.weight", loader.load_tensor(idx + "k_norm.weight")) + loader.copy_in(idx + "k_norm.bias", loader.load_tensor(idx + "k_norm.bias")) self.model.layers[i].self_attn.process_weights_after_loading() if i < cfg.first_k_dense_replace: @@ -1146,28 +1098,39 @@ def load_weights( 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([ - loader.shard(gw, 0, shard_world, shard_rank), - loader.shard(uw, 0, shard_world, shard_rank), - ], dim=0).contiguous()) + torch.cat( + [ + loader.shard(gw, 0, shard_world, shard_rank), + loader.shard(uw, 0, shard_world, shard_rank), + ], + dim=0, + ).contiguous() + ) w13_scale.data[local_idx].copy_( - torch.cat([ - loader.shard(gs, 0, shard_world, shard_rank), - loader.shard(us, 0, shard_world, shard_rank), - ], dim=0).contiguous()) + torch.cat( + [ + loader.shard(gs, 0, shard_world, shard_rank), + loader.shard(us, 0, shard_world, shard_rank), + ], + dim=0, + ).contiguous() + ) w13_offset.data[local_idx].copy_( - torch.cat([ - loader.shard(go, 0, shard_world, shard_rank), - loader.shard(uo, 0, shard_world, shard_rank), - ], dim=0).contiguous()) - w2_param.data[local_idx].copy_( - loader.shard(dw, 1, shard_world, shard_rank).contiguous()) + torch.cat( + [ + loader.shard(go, 0, shard_world, shard_rank), + loader.shard(uo, 0, shard_world, shard_rank), + ], + dim=0, + ).contiguous() + ) + 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()) - loader.copy_in(p + "mlp.gate.weight", - loader.load_tensor(p + "mlp.gate.weight")) - loader.copy_in(p + "mlp.e_score_correction_bias", - loader.load_tensor(p + "mlp.gate.e_score_correction_bias")) + loader.copy_in(p + "mlp.gate.weight", loader.load_tensor(p + "mlp.gate.weight")) + loader.copy_in( + p + "mlp.e_score_correction_bias", loader.load_tensor(p + "mlp.gate.e_score_correction_bias") + ) saved_tp = (loader.tp_size, loader.tp_rank) loader.tp_size = shard_world loader.tp_rank = shard_rank diff --git a/xllm/python/models/deepseek_v4.py b/xllm/python/models/deepseek_v4.py new file mode 100644 index 0000000000..b6a90fcfa7 --- /dev/null +++ b/xllm/python/models/deepseek_v4.py @@ -0,0 +1,1823 @@ +# 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. + +"""DeepSeek-V4 Python model (DSA sparse attention, TORCH backend). + +Structural port of the C++ ``DeepseekV4ModelImpl`` / +``DeepseekV4DecoderLayerImpl`` / ``DSAttentionImpl`` (xllm/models/llm/ +deepseek_v4.h, xllm/core/layers/deepseek_v4_decoder_layer.cpp, +xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp). Reuses the W8A8 +linear / MLP / MoE / YaRN-RoPE / weight-loader primitives from +``deepseek_v32`` and adds the DeepSeek-V4-specific pieces: + + * ``DeepseekV4Config`` -- reads the DSV4 fields (compress_ratios, window_size, + o_lora_rank, o_groups, hc_*, index_*). + * HyperConnection residual path (hc_pre / hc_post). + * ``DeepseekV4Attention`` -- q_a/kv projections + RoPE, hands q/kv to the + DSA attention backend (``backend.execute``), two-stage o_a/o_b output proj. + * ``DeepseekV4Indexer`` -- Hadamard rotation + compressor + quantized + lightning indexer. + * ``DeepseekV4ForCausalLM`` -- ``load_weights`` reusing ``W8A8WeightLoader``. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field, replace +from typing import Any + +import torch +import torch.nn as nn + +from xllm.python.attention.dsa_attention import ( + _get_layer_cache_tensor, + _scatter_by_slot, +) +from xllm.python.layers.attention import Attention +from xllm.python.layers.embedding import HiddenParallelEmbedding +from xllm.python.layers.layernorm import RMSNorm +from xllm.python.layers.linear import ColumnParallelLinear, RowParallelLinear +from xllm.python.model_executor.forward_context import ( + get_forward_context, + record_layer_event, +) +from xllm.python.models.base import PyModelBase +from xllm.python.models.deepseek_v32 import ( + DeepseekV3MLP, + DeepseekYarnRotaryEmbedding, + W8A8DynamicLinear, + W8A8WeightLoader, + _tp_rank_from_device, +) + +try: + from xllm.python import distributed +except Exception: # pragma: no cover - distributed is optional in tests + distributed = None # type: ignore[assignment] + +try: + from xllm.python import kernels +except Exception: # pragma: no cover - kernels need the compiled lib + kernels = None # type: ignore[assignment] + + +def _pick(d: dict, *keys: str, default: Any = None) -> Any: + for k in keys: + if k in d and d[k] is not None: + return d[k] + return default + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +@dataclass +class DeepseekV4Config: + """DeepSeek-V4 model config (DSV3.2 MLA fields + DSV4-specific fields).""" + + model_type: str = "deepseek_v4" + hidden_size: int = 4096 + n_layers: int = 43 + n_heads: int = 64 + head_dim: int = 512 + vocab_size: int = 129280 + rms_norm_eps: float = 1e-6 + rope_theta: float = 10000.0 + max_position_embeddings: int = 1048576 + original_max_position_embeddings: int = 65536 + rope_scaling_factor: float = 16.0 + rope_beta_fast: int = 32 + rope_beta_slow: int = 1 + rope_mscale: float = 1.0 + rope_mscale_all_dim: float = 1.0 + q_lora_rank: int = 1024 + kv_lora_rank: int = 0 + qk_nope_head_dim: int = 0 + qk_rope_head_dim: int = 64 + v_head_dim: int = 0 + # DSV4-specific + rope_head_dim: int = 64 + o_lora_rank: int = 1024 + o_groups: int = 8 + compress_ratios: list[int] = field(default_factory=list) + compress_rope_theta: float = 160000.0 + window_size: int = 128 + n_activated_experts: int = 6 + n_hash_layers: int = 3 + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + scoring_func: str = "sqrtsoftplus" + scale_fmt: str = "ue8m0" + index_head_dim: int = 128 + index_n_heads: int = 64 + index_topk: int = 512 + n_routed_experts: int = 256 + n_shared_experts: int = 1 + moe_intermediate_size: int = 2048 + swiglu_limit: float = 10.0 + first_k_dense_replace: int = 0 + moe_layer_freq: int = 1 + norm_topk_prob: bool = True + routed_scaling_factor: float = 1.5 + topk_method: str = "noaux_tc" + n_group: int = 0 + topk_group: int = 0 + tie_word_embeddings: bool = False + tp_size: int = 1 + tp_rank: int = 0 + moe_tp_size: int = 1 + moe_tp_rank: int = 0 + ep_size: int = 1 + ep_rank: int = 0 + cp_size: int = 1 + cp_rank: int = 0 + + @classmethod + def from_dict(cls, d: dict) -> DeepseekV4Config: + rs_raw = d.get("rope_scaling") + rs = rs_raw if isinstance(rs_raw, dict) else {} + + def rope_value( + model_arg: str, + nested_key: str, + default: float | int, + *legacy_keys: str, + ) -> float | int: + # PyCausalLM reflects ModelArgs into this dict. DSV4 uses factor, + # beta_fast/beta_slow and rope_scaling_attn_factor directly; older + # generic aliases may also be present with their zero defaults. + for key in (model_arg, *legacy_keys): + value = d.get(key) + if value not in (None, 0, 0.0): + return value + value = rs.get(nested_key) + return default if value in (None, 0, 0.0) else value + + n_layers = int(_pick(d, "num_hidden_layers", "n_layers", default=43)) + compress_ratios = [1 if int(ratio) <= 1 else int(ratio) for ratio in d.get("compress_ratios", [])] + if len(compress_ratios) < n_layers: + compress_ratios.extend([1] * (n_layers - len(compress_ratios))) + + return cls( + model_type=_pick(d, "model_type", default="deepseek_v4"), + hidden_size=int(_pick(d, "hidden_size", default=4096)), + n_layers=n_layers, + n_heads=int(_pick(d, "n_heads", "num_attention_heads", default=64)), + head_dim=int(_pick(d, "head_dim", default=512)), + vocab_size=int(_pick(d, "vocab_size", default=129280)), + rms_norm_eps=float(_pick(d, "rms_norm_eps", default=1e-6)), + rope_theta=float(_pick(d, "rope_theta", default=10000.0)), + max_position_embeddings=int(_pick(d, "max_position_embeddings", default=1048576)), + original_max_position_embeddings=int( + rope_value( + "rope_scaling_original_max_position_embeddings", + "original_max_position_embeddings", + 65536, + ) + ), + rope_scaling_factor=float(rope_value("factor", "factor", 16.0, "rope_scaling_factor")), + rope_beta_fast=int(rope_value("beta_fast", "beta_fast", 32, "rope_scaling_beta_fast")), + rope_beta_slow=int(rope_value("beta_slow", "beta_slow", 1, "rope_scaling_beta_slow")), + rope_mscale=float(rope_value("rope_scaling_attn_factor", "attn_factor", 1.0)), + rope_mscale_all_dim=1.0, + q_lora_rank=int(_pick(d, "q_lora_rank", default=1024)), + qk_rope_head_dim=int(_pick(d, "qk_rope_head_dim", default=64)), + rope_head_dim=int(_pick(d, "qk_rope_head_dim", default=64)), + o_lora_rank=int(_pick(d, "o_lora_rank", default=1024)), + o_groups=int(_pick(d, "o_groups", default=8)), + compress_ratios=compress_ratios, + compress_rope_theta=float(_pick(d, "compress_rope_theta", default=160000.0)), + window_size=( + int(v) if (v := _pick(d, "window_size", "sliding_window", default=128)) not in (None, -1, 0) else 128 + ), + n_activated_experts=int(_pick(d, "n_activated_experts", "num_experts_per_tok", default=6)), + n_hash_layers=int(_pick(d, "num_hash_layers", default=3)), + hc_mult=int(_pick(d, "hc_mult", default=4)), + hc_sinkhorn_iters=int(_pick(d, "hc_sinkhorn_iters", default=20)), + hc_eps=float(_pick(d, "hc_eps", default=1e-6)), + scoring_func=_pick(d, "scoring_func", default="sqrtsoftplus"), + scale_fmt=_pick(d, "scale_fmt", default="ue8m0"), + index_head_dim=int(_pick(d, "index_head_dim", default=128)), + index_n_heads=int(_pick(d, "index_n_heads", default=64)), + index_topk=int(_pick(d, "index_topk", default=512)), + n_routed_experts=int(_pick(d, "n_routed_experts", default=256)), + n_shared_experts=int(_pick(d, "n_shared_experts", default=1)), + moe_intermediate_size=int(_pick(d, "moe_intermediate_size", default=2048)), + swiglu_limit=float(_pick(d, "swiglu_limit", default=10.0)), + first_k_dense_replace=int(_pick(d, "first_k_dense_replace", default=0)), + moe_layer_freq=int(_pick(d, "moe_layer_freq", default=1)), + norm_topk_prob=bool(_pick(d, "norm_topk_prob", default=True)), + routed_scaling_factor=float(_pick(d, "routed_scaling_factor", default=1.5)), + topk_method=_pick(d, "topk_method", default="noaux_tc"), + n_group=int(_pick(d, "n_group", default=0)), + topk_group=int(_pick(d, "topk_group", default=0)), + tie_word_embeddings=bool(_pick(d, "tie_word_embeddings", default=False)), + tp_size=int(d.get("tp_size", 1)), + tp_rank=int(d.get("tp_rank", _tp_rank_from_device(d.get("device", "npu:0")))), + moe_tp_size=int(d.get("moe_tp_size", 1)), + moe_tp_rank=int(d.get("moe_tp_rank", 0)), + ep_size=int(d.get("ep_size", d.get("tp_size", 1))), + ep_rank=int(d.get("ep_rank", d.get("tp_rank", 0))), + cp_size=int(d.get("cp_size", 1)), + cp_rank=int(d.get("cp_rank", 0)), + ) + + def head_split(self) -> tuple[int, int]: + return self.n_heads // self.tp_size, 1 + + # -- aliases so DSV3.2-reused modules (MoE/MLP) read DSV4 config unchanged -- + @property + def num_experts_per_tok(self) -> int: + return self.n_activated_experts + + @property + def intermediate_size(self) -> int: + return self.moe_intermediate_size + + +# --------------------------------------------------------------------------- +# DeepSeek-V4 RoPE +# --------------------------------------------------------------------------- + + +class DeepseekV4RotaryEmbedding(nn.Module): + """Compact-cache equivalent of C++ DeepseekV4RotaryEmbedding. + + C++ keeps cache length and YaRN's old-context length as separate inputs. + The generic Python DeepseekYarnRotaryEmbedding derives cache length as + ``old_context * factor``, which cannot represent the native DSV4 call. + """ + + _cache_lock = threading.Lock() + _cache_by_descriptor: dict[tuple[Any, ...], torch.Tensor] = {} + + def __init__( + self, + rotary_dim: int, + max_position_embeddings: int, + scaling_factor: float, + theta: float, + beta_fast: int, + beta_slow: int, + old_context_len: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__() + descriptor = ( + rotary_dim, + max_position_embeddings, + scaling_factor, + theta, + beta_fast, + beta_slow, + old_context_len, + dtype, + torch.device(device), + ) + with self._cache_lock: + cache = self._cache_by_descriptor.get(descriptor) + if cache is None: + cache = self._build_cache( + rotary_dim, + max_position_embeddings, + scaling_factor, + theta, + beta_fast, + beta_slow, + old_context_len, + dtype, + device, + ) + self._cache_by_descriptor[descriptor] = cache + self.register_buffer("cos_sin_cache", cache, persistent=False) + + @staticmethod + def _build_cache( + rotary_dim: int, + max_position_embeddings: int, + scaling_factor: float, + theta: float, + beta_fast: int, + beta_slow: int, + old_context_len: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + # Match C++ rotary::apply_deepseek_yarn_rope_scaling() and + # create_cos_sin_tensor(): build the cache in CPU float32, then perform + # one final transfer/conversion to the model device and dtype. Computing + # pow/cos/sin directly on NPU can change BF16 cache entries by one ULP. + cpu = torch.device("cpu") + inv_freq = DeepseekYarnRotaryEmbedding._yarn_inv_freq( + scaling_factor, + rotary_dim, + theta, + beta_fast, + beta_slow, + old_context_len, + cpu, + ) + positions = torch.arange(max_position_embeddings, dtype=torch.float32, device=cpu) + freqs = torch.outer(positions, inv_freq) + # Keep one value per frequency. Call sites repeat_interleave to the + # same [position, rotary_dim] interleaved layout C++ stores directly. + cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1).to(device=device, dtype=dtype) + return cache.contiguous() + + +# --------------------------------------------------------------------------- +# HyperConnection +# --------------------------------------------------------------------------- + + +class DeepseekV4HyperConnection(nn.Module): + """HyperConnection residual path (hc_pre + hc_post). + + Faithful port of the C++ ``DeepseekV4DecoderLayerImpl::hc_pre``/``hc_post`` + (deepseek_v4_decoder_layer.cpp:238+), which call the registered NPU + ``hc_pre``/``hc_post`` kernels. hc_pre mixes the hc_mult parallel residual + streams into one sub-block input (via Sinkhorn); hc_post combines the + sub-block output with the residual. Weight shapes match the checkpoint: + ``hc_fn = [mix_hc, hc_dim]`` where ``mix_hc = (2+hc_mult)*hc_mult`` and + ``hc_dim = hc_mult*hidden``; ``hc_base = [mix_hc]``; ``hc_scale = [3]``. + """ + + def __init__(self, cfg: DeepseekV4Config, dtype: torch.dtype, device: torch.device) -> None: + super().__init__() + self.hc_mult = cfg.hc_mult + self.hc_eps = cfg.hc_eps + self.norm_eps = cfg.rms_norm_eps + self.sinkhorn_iters = cfg.hc_sinkhorn_iters + hidden = cfg.hidden_size + # hc_mult is a fixed model constant; NOT TP-sharded (C++ matches). + self.hc_mult_local = cfg.hc_mult + mix_hc = (2 + cfg.hc_mult) * cfg.hc_mult + hc_dim = cfg.hc_mult * hidden + # hc_fn/scale/base per sub-block. C++ registers these as float32. + for part in ("attn", "ffn"): + self.register_parameter( + f"hc_{part}_fn", + nn.Parameter(torch.empty(mix_hc, hc_dim, dtype=torch.float32, device=device)), + ) + self.register_parameter( + f"hc_{part}_base", + nn.Parameter(torch.empty(mix_hc, dtype=torch.float32, device=device)), + ) + self.register_parameter( + f"hc_{part}_scale", + nn.Parameter(torch.empty(3, dtype=torch.float32, device=device)), + ) + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Call the registered hc_pre kernel: x [T, hc_mult, hidden] -> (attn_input, post, comb).""" + from xllm.python import kernels + + return kernels.hc_pre( + x, + hc_fn, + hc_scale, + hc_base, + self.hc_mult, + self.sinkhorn_iters, + self.norm_eps, + self.hc_eps, + ) + + def hc_post( + self, + sub_out: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + """Call the registered hc_post kernel: combine sub-block output + residual. + + Faithful port of C++ (decoder_layer.cpp:258-272): when x=2D, residual=3D, + post=2D, comb=3D, unsqueeze(0) all inputs (kernel expects 3D/4D/3D/4D), + then squeeze(0) the output. + """ + from xllm.python import kernels + + if sub_out.dim() == 2 and residual.dim() == 3 and post.dim() == 2 and comb.dim() == 3: + out = kernels.hc_post( + sub_out.unsqueeze(0), + residual.unsqueeze(0), + post.unsqueeze(0), + comb.unsqueeze(0), + ) + return out.squeeze(0) + return kernels.hc_post(sub_out, residual, post, comb) + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class DeepseekV4Attention(Attention): + """DeepSeek-V4 DSA attention. + + Projects q/kv (W8A8), applies RoPE, hands q/kv to the DSA backend + (``backend.execute``), then runs the two-stage o_a/o_b output projection. + The compressor + indexer are wired through backend callbacks so the + attention forward stays a thin orchestrator matching the C++ flow. + """ + + def __init__( + self, + cfg: DeepseekV4Config, + layer_id: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + tp = cfg.tp_size + num_heads = cfg.n_heads // tp + head_dim = cfg.head_dim + scale = head_dim**-0.5 + super().__init__( + num_heads=num_heads, + num_kv_heads=1, + head_dim=head_dim, + scale=scale, + sliding_window=cfg.window_size, + layer_id=layer_id, + ) + self.cfg = cfg + self.layer_id = layer_id + self.num_heads_local = num_heads + self.head_dim = head_dim + self.rope_head_dim = cfg.qk_rope_head_dim + self.nope_head_dim = head_dim - cfg.qk_rope_head_dim + self.kv_lora_rank = cfg.q_lora_rank + # Attention sink (learnable logit bias per head), loaded from + # ``attn_sink`` in the checkpoint. Registered as float32 (the + # sparse_attn_sharedkv kernel requires DT_FLOAT sinks). + self.register_parameter( + "attn_sink", + nn.Parameter(torch.empty(num_heads, dtype=torch.float32, device=device)), + ) + # q/kv down-projections (W8A8). + # Native DSV4 keeps dynamic-W8A8 weights in checkpoint [N, K] layout + # and calls quant_matmul with transpose2=true. + self.q_a_proj = W8A8DynamicLinear( + cfg.hidden_size, + cfg.q_lora_rank, + device, + transpose_weight_after_loading=False, + ) + self.kv_proj = W8A8DynamicLinear( + cfg.hidden_size, + head_dim, + device, + transpose_weight_after_loading=False, + ) + self.q_a_layernorm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps, dtype=dtype, device=device) + self.kv_a_layernorm = RMSNorm(head_dim, cfg.rms_norm_eps, dtype=dtype, device=device) + # q up-projection (W8A8) produces [T, num_heads, head_dim]. + self.q_b_proj = W8A8DynamicLinear( + cfg.q_lora_rank, + num_heads * head_dim, + device, + transpose_weight_after_loading=False, + ) + self.register_buffer( + "q_rms_gamma", + torch.ones(head_dim, dtype=dtype, device=device), + ) + # Two-stage output projection: o_a (column) -> grouped low-rank -> o_b (row). + # o_a input is the per-group head slice = global_num_heads * head_dim / o_groups + # (uses the GLOBAL head count, not the TP-local one); C++ matches this. + assert cfg.o_groups % tp == 0 + self.n_local_groups = cfg.o_groups // tp + self.o_lora_rank = cfg.o_lora_rank + o_a_in = (cfg.n_heads * head_dim) // cfg.o_groups + # ColumnParallelLinear takes out_features_PER_PARTITION, so shard the + # full o_groups*o_lora output by tp (matches the C++ ColumnParallelLinear + # which hands the per-partition count). + self.o_a_proj = ColumnParallelLinear( + o_a_in, + (cfg.o_groups * cfg.o_lora_rank) // tp, + tp, + dtype=dtype, + device=device, + ) + self.o_b_proj = RowParallelLinear( + (cfg.o_groups * cfg.o_lora_rank) // tp, + cfg.hidden_size, + tp, + dtype=dtype, + device=device, + ) + compress_ratio = cfg.compress_ratios[layer_id] + self.indexer: DeepseekV4Indexer | None = ( + DeepseekV4Indexer(cfg, dtype, device) if compress_ratio == 4 and cfg.index_topk > 0 else None + ) + # Cmp_kv compressor (separate from the indexer compressor). C++ DSA + # attention has its own CompressorImpl with head_dim_=512 (attention + # head_dim), distinct from the indexer's head_dim_=128. Weights loaded + # from attn.compressor.* (not attn.indexer.compressor.*). + # wkv/wgate out = coff * head_dim (coff=2 for C4, 1 for C128); + # norm=head_dim; ape=[4, coff*head_dim]. + cmp_hd = head_dim # attention head_dim=512, NOT index_head_dim=128 + if compress_ratio > 1: + cmp_coff = 2 if compress_ratio == 4 else 1 + cmp_out = cmp_coff * cmp_hd + self.cmp_wkv = nn.Linear( + cfg.hidden_size, + cmp_out, + bias=False, + dtype=torch.float32, + device=device, + ) + self.cmp_wgate = nn.Linear( + cfg.hidden_size, + cmp_out, + bias=False, + dtype=torch.float32, + device=device, + ) + self.cmp_ape = nn.Parameter(torch.empty(compress_ratio, cmp_out, dtype=torch.float32, device=device)) + self.cmp_norm = RMSNorm(cmp_hd, cfg.rms_norm_eps, dtype=torch.float32, device=device) + + def process_weights_after_loading(self) -> None: + for m in (self.q_a_proj, self.kv_proj, self.q_b_proj): + m.process_weights_after_loading() + # Keep o_b in checkpoint [N, K] layout. Native DSAttention sends this + # unquantized RowParallelLinear through F.linear(input, weight); the + # generic NPU preparation transposes it to FRACTAL_NZ and selects a + # different matmul accumulation path. + if hasattr(self.o_a_proj, "process_weights_after_loading"): + self.o_a_proj.process_weights_after_loading() + if self.indexer is not None: + self.indexer.process_weights_after_loading() + + def forward( + self, + hidden: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden.shape[0] + backend = get_forward_context().attention_backend + metadata = get_forward_context().metadata + dsa = getattr(metadata, "dsa_metadata", None) + kv_hidden = hidden + + # q/kv down + up + RoPE (matches run_dsv4_preprocess_fallback). + # W8A8 path (C++ deepseek_sparse_attention.cpp:396-413): q_a_proj does + # dynamic_quant internally; then rms_norm_dynamic_quant fuses the q_a + # RMSNorm with a second dynamic quant, producing qr (int8) + + # qr_pertoken_scale; q_b_proj consumes that pre-quantized qr (no + # re-quant). The indexer reuses the same qr + qr_pertoken_scale in + # build_query, so stash them on the backend for _run_indexer. + q_a = self.q_a_proj(hidden) + from xllm.python import kernels as _k + + qr, qr_pertoken_scale = _k.rms_norm_dynamic_quant(q_a, self.q_a_layernorm.weight, self.cfg.rms_norm_eps) + q = self.q_b_proj.forward_quant(qr, qr_pertoken_scale).view(num_tokens, self.num_heads_local, self.head_dim) + q = _k.rms_norm(q, self.q_rms_gamma, self.cfg.rms_norm_eps) + + cos_sin = cos_sin_cache.index_select(0, positions.long()) + half = cos_sin.size(-1) // 2 + cos = cos_sin[..., :half].repeat_interleave(2, dim=-1).contiguous() + sin = cos_sin[..., half:].repeat_interleave(2, dim=-1).contiguous() + _k.npu_inplace_partial_rotary_mul(q, cos, sin, self.nope_head_dim, self.rope_head_dim) + + kv = self.kv_proj(kv_hidden) + # kv_proj outputs head_dim = nope_head_dim + rope_head_dim; layernorm + # the whole thing then split for RoPE (matches C++ run_dsv4_preprocess). + kv = self.kv_a_layernorm(kv) + kv_tensor = kv.view(kv_hidden.shape[0], 1, self.head_dim) + kv_cos, kv_sin = cos, sin + _k.npu_inplace_partial_rotary_mul( + kv_tensor, + kv_cos, + kv_sin, + self.nope_head_dim, + self.rope_head_dim, + ) + + # Attach the compressor/indexer callbacks so the backend can invoke them. + if self.indexer is not None and hasattr(backend, "attach_indexer"): + backend.attach_indexer(self._run_indexer) + if hasattr(backend, "attach_compressor"): + backend.attach_compressor(self._run_compressor) + + # Pass hidden to backend so compressor/indexer can access it. + backend._current_hidden = hidden + backend._current_kv_hidden = kv_hidden + # Stash the W8A8 pre-quantized query (int8 qr + per-token scale) for the + # indexer build_query path (mirrors C++ select_qli's qr/qr_pertoken_scale). + backend._current_qr = qr + backend._current_qr_pertoken_scale = qr_pertoken_scale + attn_out = backend.execute(q, kv_tensor, kv_tensor, self) + # Native DSA rotates the attention output back before o_a/o_b. + _k.npu_inplace_partial_rotary_mul( + attn_out, + cos, + sin, + self.nope_head_dim, + self.rope_head_dim, + inverse=True, + ) + # Two-stage output projection (o_a -> grouped -> o_b). + num_tokens = attn_out.size(0) + out = attn_out.view(num_tokens, self.n_local_groups, -1) + # Match C++ DSAttentionImpl exactly. A flattened F.linear is + # mathematically equivalent but selects a different NPU accumulation + # path and produces layer-by-layer BF16 drift. + wo_a = self.o_a_proj.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o_low = torch.einsum("tgd,grd->tgr", out, wo_a) + o = self.o_b_proj(o_low.reshape(num_tokens, -1)) + # o_b_proj is RowParallelLinear with reduce_results=True (default), + # which internally calls tp_all_reduce. Do NOT call tp_all_reduce again + # here — that would be a duplicate collective (the expert analysis + # identified this as a cause of HCCL deadlock / 507015). + return o + + def _run_compressor(self, layer_id, layer_cache, dsa, mapping, cmp_block_table, compress_ratio): + """Cmp_kv compressor callback (attention-level, head_dim=512). + + Mirrors C++ DSAttentionImpl's compressor_->forward + (deepseek_sparse_attention.cpp:840-852): uses the attention's own + CompressorImpl with head_dim_=512, NOT the indexer's head_dim_=128. + Weights from attn.compressor.* (loaded separately from indexer's). + """ + if kernels is None: + return None + if not hasattr(self, "cmp_wkv"): + return None + backend = get_forward_context().attention_backend + hidden = getattr(backend, "_current_kv_hidden", None) + if hidden is None: + return None + kv_state = layer_cache.compress_kv_state + score_state = layer_cache.compress_score_state + if kv_state is None or score_state is None: + return None + # Select per-ratio compressed RoPE table (c4 for ratio=4, c128 for 128). + if compress_ratio == 4: + cos_table = dsa.c4_cos + sin_table = dsa.c4_sin + elif compress_ratio == 128: + cos_table = dsa.c128_cos + sin_table = dsa.c128_sin + else: + cos_table = dsa.cos_table + sin_table = dsa.sin_table + if cos_table is None or sin_table is None: + return None + kv_block_table = ( + _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.kv_state_cache_idx) + if dsa.block_tables + else None + ) + score_block_table = ( + _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.score_state_cache_idx) + if dsa.block_tables + else None + ) + coff = 2 if compress_ratio == 4 else 1 + rope_head_dim = self.cfg.qk_rope_head_dim + sin_view = sin_table.reshape(-1, sin_table.size(-1)) if sin_table.dim() > 2 else sin_table + cos_view = cos_table.reshape(-1, cos_table.size(-1)) if cos_table.dim() > 2 else cos_table + # Interleaved RoPE: repeat_interleave half-dim to full rope_head_dim. + if sin_view.size(-1) * 2 == rope_head_dim: + sin_view = sin_view.repeat_interleave(2, dim=-1).contiguous() + cos_view = cos_view.repeat_interleave(2, dim=-1).contiguous() + # Keep host-side metadata on CPU (C++ passes cu_seqlens/start_pos/ + # block_table as CPU tensors; aclnnCompressor tiling reads them from + # host. Moving them to NPU breaks the contract → 507015 aicore). + seq_q = dsa.actual_seq_lengths_query.contiguous() + start_pos = dsa.start_pos.contiguous() if dsa.start_pos.numel() > 0 else None + bf16 = torch.bfloat16 + # Dump ACTUAL kernel args (post-cast) with stride+format for C++ comparison. + x_arg = hidden.to(bf16).contiguous() + wkv_arg = self.cmp_wkv.weight.to(bf16).contiguous() + wgate_arg = self.cmp_wgate.weight.to(bf16).contiguous() + norm_arg = self.cmp_norm.weight.to(bf16).contiguous() + sin_arg = sin_view.to(bf16).to(hidden.device).contiguous() + cos_arg = cos_view.to(bf16).to(hidden.device).contiguous() + compressed_kv, _, _, _, _ = kernels.compressor( + x=x_arg, + wkv=wkv_arg, + wgate=wgate_arg, + kv_state=kv_state, + score_state=score_state, + ape=self.cmp_ape, + norm_weight=norm_arg, + rope_sin=sin_arg, + rope_cos=cos_arg, + kv_block_table=kv_block_table, + score_block_table=score_block_table, + cu_seqlens=seq_q, + seqused=None, + start_pos=start_pos, + rope_head_dim=rope_head_dim, + cmp_ratio=compress_ratio, + coff=coff, + norm_eps=self.cfg.rms_norm_eps, + rotary_mode=2, + enable_grad=False, + ) + return compressed_kv + + def _run_indexer(self, layer_id, layer_cache, dsa, mapping, q): + """Indexer callback: returns top-k compressed block indices.""" + if self.indexer is not None: + backend = get_forward_context().attention_backend + qr = getattr(backend, "_current_qr", None) + qr_pertoken_scale = getattr(backend, "_current_qr_pertoken_scale", None) + hidden = getattr(backend, "_current_hidden", None) + kv_hidden = getattr(backend, "_current_kv_hidden", hidden) + return self.indexer.select_qli_dsv4( + layer_id, + layer_cache, + dsa, + mapping, + q, + qr, + qr_pertoken_scale, + hidden, + kv_hidden, + ) + return None + + +# --------------------------------------------------------------------------- +# Indexer + Compressor +# --------------------------------------------------------------------------- + + +class DeepseekV4Indexer(nn.Module): + """DeepSeek-V4 indexer: Hadamard rotation + compressor + quant lightning. + + Faithful to the C++ ``DeepseekV4IndexerImpl`` (deepseek_v4_indexer.cpp): + Hadamard-rotates the compressed key, runs the NSA compressor, scatters the + compressed key into the paged index cache, then runs the quantized lightning + indexer to pick top-k compressed blocks. + """ + + def __init__(self, cfg: DeepseekV4Config, dtype: torch.dtype, device: torch.device) -> None: + super().__init__() + self.cfg = cfg + self.n_head = cfg.index_n_heads + self.head_dim = cfg.index_head_dim + self.rope_dim = cfg.qk_rope_head_dim + self.topk = cfg.index_topk + self.dtype = dtype + self.hadamard_scale = cfg.index_head_dim**-0.5 if cfg.index_head_dim else 1.0 + # indexer q projection + scoring weights + compressor (K projection is + # done by the compressor's wkv, so there is no separate wk/k_norm -- + # matches C++ DeepseekV4IndexerImpl). + self.wq_b = W8A8DynamicLinear( + cfg.q_lora_rank, + self.n_head * self.head_dim, + device, + transpose_weight_after_loading=False, + ) + self.weights_proj = nn.Linear(cfg.hidden_size, self.n_head, bias=False, dtype=dtype, device=device) + # Compressor: wkv (fused wk+wv, unquantized f32) + wgate + ape + norm. + # wkv out = 2*head_dim (cat of wk, wv); ape = [4, 2*head_dim]. + cmp_out = 2 * self.head_dim + self.compressor_wkv = nn.Linear(cfg.hidden_size, cmp_out, bias=False, dtype=torch.float32, device=device) + self.compressor_wgate = nn.Linear(cfg.hidden_size, cmp_out, bias=False, dtype=torch.float32, device=device) + self.compressor_ape = nn.Parameter(torch.empty(4, cmp_out, dtype=torch.float32, device=device)) + self.compressor_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps, dtype=torch.float32, device=device) + + def process_weights_after_loading(self) -> None: + self.wq_b.process_weights_after_loading() + + def select_qli_dsv4( + self, + layer_id, + layer_cache, + dsa, + mapping, + q, + qr, + qr_pertoken_scale, + hidden, + kv_hidden=None, + ) -> torch.Tensor: + """Quantized lightning indexer: returns top-k compressed block indices. + + Faithful port of C++ ``DeepseekV4IndexerImpl::select_qli`` + (deepseek_v4_indexer.cpp:393-555). The C++ path does NOT just read a + pre-filled index cache -- it rebuilds it every call: + build_query(qr, qr_pertoken_scale) -> q; partial RoPE + Hadamard(q); + compress_kv(hidden) -> kv; Hadamard(kv); dynamic_quant_int8(kv) -> + kv_quant + kv_scale; scatter kv_quant -> index_cache, kv_scale -> + indexer_scale (via slot_mapping); then quant_lightning_indexer. + Python mirrors that, otherwise index_cache/indexer_scale hold stale or + uninitialized data and the returned topk addresses out-of-range blocks + (507015 aicore in sparse_attn_sharedkv). + """ + if kernels is None or layer_cache.index is None: + return torch.empty(0, dtype=torch.int32) + index_cache = layer_cache.index + device = index_cache.device + # --- build_query (C++ 310-335): wq_b W8A8 matmul over pre-quantized qr. --- + if qr is not None and qr_pertoken_scale is not None: + q_idx = self.wq_b.forward_quant(qr, qr_pertoken_scale).view(-1, self.n_head, self.head_dim) + else: + q_idx = self.wq_b(qr).view(-1, self.n_head, self.head_dim) + # --- partial RoPE on q (C++ 417-422): apply_partial_rope over + # [rope_start_dim:rope_start_dim+rope_head_dim]. Uses the DEFAULT RoPE + # table (cos/sin) indexed by positions, 2D [M, rope_dim], NOT the + # compressed c4 table. Mirrors C++ apply_partial_rope -> + # npu_inplace_partial_rotary_mul (deepseek_sparse_attention.cpp:151-190). + rope_start_dim = max(self.head_dim - self.rope_dim, 0) + cos_table = dsa.cos_table + sin_table = dsa.sin_table + if cos_table is not None and sin_table is not None and self.rope_dim > 0: + cos_v = cos_table.reshape(-1, cos_table.size(-1)) if cos_table.dim() > 2 else cos_table + sin_v = sin_table.reshape(-1, sin_table.size(-1)) if sin_table.dim() > 2 else sin_table + # Per-token cos/sin indexed by positions: 2D [M, rope_dim/2] (Python + # DeepseekYarnRotaryEmbedding stores half-dim cos/sin, NOT interleaved). + pos = dsa.input_positions.to(device).reshape(-1).long() + cos_sel = cos_v.index_select(0, pos).to(q_idx.dtype) # [M, rope_dim/2] + sin_sel = sin_v.index_select(0, pos).to(q_idx.dtype) + # npu_inplace_partial_rotary_mul (interleave mode) expects cos/sin + # [M, rope_dim] in C++ interleaved format: freqs.repeat_interleave(2) + # (rotary_embedding_util.cpp:135-137). The half-dim cos/sin must be + # repeat_interleave'd to full rope_dim. Local fix only -- do NOT change + # the global DeepseekYarnRotaryEmbedding cache (the attention main path + # uses _interleave_rope_with which consumes half-dim). + if cos_sel.size(-1) * 2 == self.rope_dim: + cos_sel = cos_sel.repeat_interleave(2, dim=-1).contiguous() + sin_sel = sin_sel.repeat_interleave(2, dim=-1).contiguous() + elif cos_sel.size(-1) != self.rope_dim: + raise RuntimeError(f"QRoPE cos/sin dim mismatch: cos={cos_sel.shape}, rope_dim={self.rope_dim}") + # In-place partial RoPE: modifies q_idx[...rope_start_dim:rope_dim] + # via aclnnInplacePartialRotaryMul (interleave mode). + from xllm.python import kernels as _pk + + _pk.npu_inplace_partial_rotary_mul(q_idx, cos_sel, sin_sel, rope_start_dim, self.rope_dim) + # --- Hadamard rotation on q (C++ 423-424). --- + hadamard = self._get_hadamard(device) + q_idx = _rotate_hadamard(q_idx, hadamard, self.hadamard_scale) + # --- build_weights(hidden) (C++ 337-340, 456). --- + softmax_mul = (self.head_dim**-0.5) * (self.n_head**-0.5) + weights = self.weights_proj(hidden) * softmax_mul + # --- Rebuild index cache: compress_kv -> Hadamard -> quant -> scatter. --- + kv = self._indexer_compress_kv( + kv_hidden if kv_hidden is not None else hidden, + layer_cache, + dsa, + mapping, + layer_id, + ) + if kv is not None and kv.numel() > 0: + kv = _rotate_hadamard(kv, hadamard, self.hadamard_scale) + kv_quant, kv_scale = kernels.dynamic_quant(kv) + kv_scale = kv_scale.unsqueeze(-1).to(torch.float16) + # Scatter kv_quant -> index_cache, kv_scale -> indexer_scale, by slot. + slot = _get_layer_cache_tensor(dsa.slot_mappings, layer_id, mapping.index_cache_idx) + if slot is not None and slot.numel() > 0: + _scatter_by_slot(index_cache, slot, kv_quant) + if layer_cache.indexer_scale is not None: + _scatter_by_slot(layer_cache.indexer_scale, slot, kv_scale) + # --- dynamic_quant_int8(q) -> q_quant (int8), q_scale (float16). --- + q_quant, q_scale = kernels.dynamic_quant(q_idx) + q_scale = q_scale.to(torch.float16) + # key_dequant_scale = indexer_scale (written above), else ones fallback. + key_dequant_scale = layer_cache.indexer_scale + if key_dequant_scale is None or key_dequant_scale.numel() == 0: + scale_sizes = list(index_cache.shape) + scale_sizes[-1] = 1 + key_dequant_scale = torch.ones(scale_sizes, dtype=torch.float16, device=device) + block_table = _layer_tensor(dsa.block_tables, layer_id, mapping.index_cache_idx) + query_seq_lens = dsa.actual_seq_lengths_query + if query_seq_lens.dim() > 0 and query_seq_lens.size(0) > 1: + query_seq_lens = query_seq_lens[1:] + key_seq_lens = dsa.actual_seq_lengths_kv + qli_metadata = dsa.qli_metadata + # C++ packs the current forward's DSA metadata onto the runtime device + # before building and invoking QLI. Python follows the same ownership + # contract in DsaAttentionBackend.prepare_dsa_metadata_for_forward(); + # do not create per-call copies here because they hide lifecycle bugs. + for name, tensor in ( + ("query_seq_lens", query_seq_lens), + ("key_seq_lens", key_seq_lens), + ("block_table", block_table), + ("qli_metadata", qli_metadata), + ): + if tensor is None or tensor.numel() == 0: + raise RuntimeError(f"QLI {name} must be defined and non-empty") + if tensor.device != device: + raise RuntimeError(f"QLI {name} must be on {device}, got {tensor.device}") + topk, _ = kernels.quant_lightning_indexer( + query=q_quant, + key=index_cache, + weights=weights.to(torch.float16), + query_dequant_scale=q_scale, + key_dequant_scale=key_dequant_scale.to(torch.float16), + query_quant_mode=0, + key_quant_mode=0, + actual_seq_lengths_query=query_seq_lens, + actual_seq_lengths_key=key_seq_lens, + block_table=block_table, + metadata=qli_metadata, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.topk, + sparse_mode=3, + pre_tokens=2**63 - 1, + next_tokens=2**63 - 1, + cmp_ratio=4, + return_value=False, + ) + return topk + + def _get_hadamard(self, device: torch.device) -> torch.Tensor: + """Build (or cache) the Sylvester Hadamard matrix for head_dim. + + Mirrors C++ ``create_hadamard_matrix`` (index_head_dim_padded = next + pow2 >= head_dim, normalize=False). Cached on the module. + """ + cached = getattr(self, "_hadamard_matrix", None) + if cached is not None and cached.device == device: + return cached + n = 1 + while n < self.head_dim: + n <<= 1 + mat = torch.ones((1, 1), dtype=self.dtype, device=device) + m = 1 + while m < n: + top = torch.cat([mat, mat], 1) + bottom = torch.cat([mat, -mat], 1) + mat = torch.cat([top, bottom], 0) + m <<= 1 + self._hadamard_matrix = mat + return mat + + def _indexer_compress_kv(self, hidden, layer_cache, dsa, mapping, layer_id): + """Run the compressor against the index-cache states (not cmp_kv). + + Mirrors C++ ``select_qli``'s ``compress_kv(kv_source, ..., + &indexer_states, &indexer_block_tables, c4_cos, c4_sin, ...)``. Uses + ``compress_index_kv_state`` / ``compress_index_score_state`` + the + index-cache block tables, distinct from the cmp_kv compressor callback. + """ + kv_state = layer_cache.compress_index_kv_state + score_state = layer_cache.compress_index_score_state + if kv_state is None or score_state is None: + return None + # Indexer compressor always runs on C4 layers (backend execute guards + # compress_ratio==4), so use the c4 compressed RoPE table, not the full + # base table (compressor.cpp:48-60 derives cmp_s = rope_sin.size(0)). + cos_table = dsa.c4_cos + sin_table = dsa.c4_sin + if cos_table is None or sin_table is None: + return None + kv_block_table = ( + _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.index_kv_state_cache_idx) + if dsa.block_tables + else None + ) + score_block_table = ( + _get_layer_cache_tensor(dsa.block_tables, layer_id, mapping.index_score_state_cache_idx) + if dsa.block_tables + else None + ) + # Indexer only runs on C4 layers (backend execute: compress_ratio==4), + # so cmp_ratio is always 4 here -- not a hardcode bug, unlike cmp_kv. + cmp_ratio = 4 + coff = 2 if cmp_ratio == 4 else 1 + rope_head_dim = self.cfg.qk_rope_head_dim + sin_view = sin_table.reshape(-1, sin_table.size(-1)) if sin_table.dim() > 2 else sin_table + cos_view = cos_table.reshape(-1, cos_table.size(-1)) if cos_table.dim() > 2 else cos_table + # Compressor RoPE is interleaved; Python compress_rotary_c4 stores + # half-dim [M, rope_dim/2]. repeat_interleave to [M, rope_dim]. + if sin_view.size(-1) * 2 == rope_head_dim: + sin_view = sin_view.repeat_interleave(2, dim=-1).contiguous() + cos_view = cos_view.repeat_interleave(2, dim=-1).contiguous() + # Keep host-side metadata on CPU (same fix as cmp_kv compressor). + seq_q = dsa.actual_seq_lengths_query.contiguous() + start_pos = dsa.start_pos.contiguous() if dsa.start_pos.numel() > 0 else None + # aclnnCompressor dtype contract (from the supported-list error): + # x/wkv/wgate/normWeight/ropeSin/ropeCos/cmpKvOut = bf16 (or fp16); + # kv_state/score_state/ape = f32. Cast the bf16-group inputs. + bf16 = torch.bfloat16 + compressed_kv, _, _, _, _ = kernels.compressor( + x=hidden.to(bf16), + wkv=self.compressor_wkv.weight.to(bf16), + wgate=self.compressor_wgate.weight.to(bf16), + kv_state=kv_state, + score_state=score_state, + ape=self.compressor_ape, + norm_weight=self.compressor_norm.weight.to(bf16), + rope_sin=sin_view.to(bf16).to(hidden.device), + rope_cos=cos_view.to(bf16).to(hidden.device), + kv_block_table=kv_block_table, + score_block_table=score_block_table, + cu_seqlens=seq_q, + seqused=None, + start_pos=start_pos, + rope_head_dim=rope_head_dim, + cmp_ratio=cmp_ratio, + coff=coff, + norm_eps=self.cfg.rms_norm_eps, + rotary_mode=2, + enable_grad=False, + ) + return compressed_kv + + +def _layer_tensor(block_tables, layer_id: int, cache_idx: int): + """Fetch a per-layer block table by cache index. + + ``block_tables`` is ``DsaMetadata.block_tables``: ``[n_layers][n_caches]`` + with the same underlying tensor shared across caches in one group. Mirrors + C++ ``get_layer_cache_tensor`` (deepseek_sparse_attention.cpp:80). + """ + if layer_id < 0 or layer_id >= len(block_tables) or cache_idx < 0 or cache_idx >= len(block_tables[layer_id]): + return None + tensor = block_tables[layer_id][cache_idx] + return tensor if tensor.numel() > 0 else None + + +def _rotate_hadamard(x: torch.Tensor, hadamard: torch.Tensor, scale: float) -> torch.Tensor: + """Apply the Hadamard transform to the last dim of ``x``. + + Faithful port of C++ ``rotate_activation_with_hadamard`` -> + ``hadamard_transform_ref`` (deepseek_v4_indexer.cpp:61-87): pad the last + dim to ``hadamard.size(0)`` (next pow2), matmul, slice back, scale. + """ + if hadamard is None or hadamard.numel() == 0: + return x + dim = x.size(-1) + x2d = x.reshape(-1, dim) + if x2d.dtype != hadamard.dtype: + raise RuntimeError(f"Hadamard dtype must match input: {hadamard.dtype} != {x2d.dtype}") + dim_padded = hadamard.size(0) + if dim != dim_padded: + x2d = torch.nn.functional.pad(x2d, (0, dim_padded - dim)) + out = torch.nn.functional.linear(x2d, hadamard) + out = out[:, :dim] + out = out.reshape(x.shape) + if scale != 1.0: + out = out * scale + return out + + +# --------------------------------------------------------------------------- +# DSV4 MoE with hash routing +# --------------------------------------------------------------------------- + + +class DeepseekV4MoE(nn.Module): + """DeepSeek-V4 MoE with hash routing. + + Uses ``moe_gating_top_k_hash`` for routing (hash layers) or bias-based + routing (non-hash layers), then ``grouped_moe_with_selected_experts`` for + expert computation, plus shared experts. Mirrors C++ DeepseekV4GateImpl + + FusedMoEImpl::forward_with_selected_experts. + """ + + def __init__(self, cfg: DeepseekV4Config, layer_id: int, dtype: torch.dtype, device: torch.device) -> None: + super().__init__() + self.cfg = cfg + self.layer_id = layer_id + self.topk = cfg.n_activated_experts + self.num_total_experts = cfg.n_routed_experts + self.routed_scaling = cfg.routed_scaling_factor + self.n_hash_layers = cfg.n_hash_layers + self.hash_layer = 0 <= layer_id < cfg.n_hash_layers + self.scoring_func = cfg.scoring_func + # EP: each rank holds num_experts_per_rank experts (not all). + # Mirrors C++ FusedMoEImpl (fused_moe.cpp:420-421). + # Use ep_size/ep_rank from config dict (set by PyCausalLM::build_config_dict). + ep_size = cfg.ep_size if cfg.ep_size > 0 else cfg.tp_size + ep_rank = cfg.ep_rank if cfg.ep_size > 0 else cfg.tp_rank + # Attention TP and MoE TP are independent under orthogonal CP. C++ + # exposes both process groups; deriving MoE TP from attention TP makes + # cp=2, ep=8 fail at construction time (attention TP=4, MoE TP=1). + self.moe_tp_size = cfg.moe_tp_size + self.moe_tp_rank = cfg.moe_tp_rank + self.num_experts_per_rank = self.num_total_experts // ep_size + self.start_expert_id = ep_rank * self.num_experts_per_rank + inter_local = cfg.moe_intermediate_size // self.moe_tp_size + self.inter_local = inter_local + + # Gate weight [n_total_experts, hidden] float32 (replicated, not sharded). + self.gate = nn.Linear(cfg.hidden_size, cfg.n_routed_experts, bias=False, dtype=torch.float32, device=device) + # Hash table for hash layers [vocab, topk] int32 (None for non-hash). + if self.hash_layer: + self.tid2eid = nn.Parameter( + torch.empty( + cfg.vocab_size, + cfg.n_activated_experts, + dtype=torch.int32, + device=device, + ), + requires_grad=False, + ) + else: + self.e_score_correction_bias = nn.Parameter( + torch.empty( + cfg.n_routed_experts, + dtype=torch.float32, + device=device, + ), + requires_grad=False, + ) + + # Expert weights — EP sharded: each rank holds num_experts_per_rank experts + # (not all n_routed_experts). Mirrors C++ FusedMoEImpl (fused_moe.cpp:605-624). + nepr = self.num_experts_per_rank + self.experts_w13 = nn.Parameter( + torch.empty(nepr, 2 * inter_local, cfg.hidden_size, dtype=torch.int8, device=device), + requires_grad=False, + ) + self.experts_w2 = nn.Parameter( + torch.empty(nepr, cfg.hidden_size, inter_local, dtype=torch.int8, device=device), + requires_grad=False, + ) + self.register_buffer( + "experts_w13_scale", torch.empty(nepr, 2 * inter_local, 1, dtype=torch.float32, device=device) + ) + self.register_buffer( + "experts_w13_offset", torch.zeros(nepr, 2 * inter_local, 1, dtype=torch.float32, device=device) + ) + self.register_buffer( + "experts_w2_scale", torch.empty(nepr, cfg.hidden_size, 1, dtype=torch.float32, device=device) + ) + self.register_buffer( + "experts_w2_offset", torch.zeros(nepr, cfg.hidden_size, 1, dtype=torch.float32, device=device) + ) + + # Shared expert uses the orthogonal MoE TP group, matching C++ + # FusedMoEImpl. skip_tp_reduce keeps collective ordering in this class. + shared_cfg = replace(cfg, tp_size=self.moe_tp_size, tp_rank=self.moe_tp_rank) + self.shared_experts = DeepseekV3MLP( + shared_cfg, + cfg.moe_intermediate_size * cfg.n_shared_experts, + dtype, + device, + skip_tp_reduce=True, + swiglu_limit=cfg.swiglu_limit, + ) + + def process_weights_after_loading(self) -> None: + # Transpose [expert, out, in] -> [expert, in, out] (matching C++ + # ensure_group_gemm_weight_layout). NO NZ — C++ forward_expert path + # does not call maybe_trans_nz; op-plugin handles format internally. + self.experts_w13.data = self.experts_w13.data.transpose(1, 2).contiguous() + self.experts_w2.data = self.experts_w2.data.transpose(1, 2).contiguous() + self.experts_w13_scale.data = self.experts_w13_scale.data.squeeze(-1).contiguous() + self.experts_w2_scale.data = self.experts_w2_scale.data.squeeze(-1).contiguous() + self.shared_experts.gate_up_proj.process_weights_after_loading() + self.shared_experts.down_proj.process_weights_after_loading() + + def forward(self, hidden: torch.Tensor, input_ids: torch.Tensor | None = None) -> torch.Tensor: + from xllm.python import kernels + + # Prepare input_ids: reshape to 1D + move to hidden's device (C++ :202-216). + gate_input_ids = None + if input_ids is not None and input_ids.numel() > 0: + flat_ids = input_ids.reshape(-1).to(hidden.device) + token_count = flat_ids.size(0) + hidden_rows = hidden.size(0) + if token_count == hidden_rows: + gate_input_ids = flat_ids + elif token_count > 0 and hidden_rows % token_count == 0: + repeat_factor = hidden_rows // token_count + gate_input_ids = flat_ids.unsqueeze(1).repeat(1, repeat_factor).reshape(hidden_rows) + + # 1) Gate: compute logits + moe_gating_top_k_hash. + gate_input = hidden.to(torch.float32) + logits = self.gate(gate_input) + norm_type = {"softmax": 0, "sigmoid": 1, "sqrtsoftplus": 2}.get(self.scoring_func, 2) + renorm = 0 if norm_type == 2 else 1 + + if self.hash_layer and hasattr(self, "tid2eid") and gate_input_ids is not None: + topk_weights, topk_idx, _ = kernels.moe_gating_top_k_hash( + x=logits, + k=self.topk, + bias=None, + input_ids=gate_input_ids, + tid2eid=self.tid2eid, + k_group=1, + group_count=1, + routed_scaling_factor=self.routed_scaling, + eps=1e-20, + group_select_mode=1, + renorm=renorm, + norm_type=norm_type, + out_flag=False, + ) + else: + bias = getattr(self, "e_score_correction_bias", None) + topk_weights, topk_idx, _ = kernels.moe_gating_top_k_hash( + x=logits, + k=self.topk, + bias=bias, + input_ids=None, + tid2eid=None, + k_group=1, + group_count=1, + routed_scaling_factor=self.routed_scaling, + eps=1e-20, + group_select_mode=1, + renorm=renorm, + norm_type=norm_type, + out_flag=False, + ) + + # 2) EP: zero out non-local expert weights (C++ fused_moe.cpp:843-850). + ep_size = self.cfg.ep_size if self.cfg.ep_size > 0 else self.cfg.tp_size + if ep_size > 1: + local_mask = (topk_idx >= self.start_expert_id) & ( + topk_idx < self.start_expert_id + self.num_experts_per_rank + ) + topk_weights = topk_weights * local_mask.to(topk_weights.dtype) + + # 3) Expert computation with pre-selected routing (EP-sharded). + routed_out = kernels.grouped_moe_with_selected_experts( + hidden, + topk_weights, + topk_idx.to(torch.int32), + self.experts_w13, + self.experts_w2, + self.experts_w13_scale, + self.experts_w2_scale, + self.experts_w13_offset, + self.experts_w2_offset, + self.num_total_experts, + self.start_expert_id, + self.num_experts_per_rank, + self.cfg.swiglu_limit, + ) + # 4) Shared experts + C++-ordered TP/EP reductions. + shared_out = self.shared_experts(hidden) + + return self._reduce_moe_outputs(routed_out, shared_out) + + def _reduce_moe_outputs(self, routed_out: torch.Tensor, shared_out: torch.Tensor) -> torch.Tensor: + """Reduce routed/shared results in the C++ ``FusedMoEImpl`` order. + + With both MoE TP and EP enabled, routed and shared outputs are partial + on different dimensions. C++ reduces each partial over MoE TP first, + reduces routed over MoE EP, then reduces shared over MoE TP before + adding the two results (fused_moe.cpp:2158-2181). Keeping this order + also keeps every rank in both process groups in the same collective + sequence. + """ + ep_size = self.cfg.ep_size if self.cfg.ep_size > 0 else self.cfg.tp_size + if distributed is None and (ep_size > 1 or self.moe_tp_size > 1): + raise RuntimeError("Python distributed collectives are unavailable") + + if ep_size > 1: + if self.moe_tp_size > 1: + distributed.moe_tp_all_reduce(routed_out) + distributed.moe_ep_all_reduce(routed_out) + if self.moe_tp_size > 1: + distributed.moe_tp_all_reduce(shared_out) + return routed_out + shared_out + + out = routed_out + shared_out + if self.moe_tp_size > 1: + # EP1: routed and shared partials can be combined before one + # reduction, matching C++'s reduce(a + b) fast path. + distributed.moe_tp_all_reduce(out) + return out + + +# --------------------------------------------------------------------------- +# Decoder layer + model +# --------------------------------------------------------------------------- + + +class DeepseekV4DecoderLayer(nn.Module): + """DeepSeek-V4 decoder layer: HyperConnection(attn) + HyperConnection(ffn).""" + + def __init__( + self, + cfg: DeepseekV4Config, + layer_id: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hc = DeepseekV4HyperConnection(cfg, dtype, device) + self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) + self.self_attn = DeepseekV4Attention(cfg, layer_id, dtype, device) + self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) + # Dense vs MoE by first_k_dense_replace / moe_layer_freq. + is_dense = (layer_id < cfg.first_k_dense_replace) or ( + cfg.moe_layer_freq > 1 and layer_id % cfg.moe_layer_freq != 0 + ) + if is_dense: + self.mlp = DeepseekV3MLP( + cfg, + cfg.moe_intermediate_size, + dtype, + device, + swiglu_limit=cfg.swiglu_limit, + ) + else: + self.mlp = DeepseekV4MoE(cfg, layer_id, dtype, device) + + def forward( + self, + hidden: torch.Tensor, + residual: torch.Tensor | None, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Match DeepseekV4DecoderLayerImpl::forward exactly: HyperConnection + # selects the 2D sub-input first, then RMSNorm is applied to that input. + residual_attn = hidden + attn_input, post_attn, comb_attn = self.hc.hc_pre( + hidden, + self.hc.hc_attn_fn, + self.hc.hc_attn_scale, + self.hc.hc_attn_base, + ) + attn_input = self.input_layernorm(attn_input) + attn_output = self.self_attn(attn_input, positions, cos_sin_cache) + hidden = self.hc.hc_post(attn_output, residual_attn, post_attn, comb_attn) + + residual_ffn = hidden + ffn_input, post_ffn, comb_ffn = self.hc.hc_pre( + hidden, + self.hc.hc_ffn_fn, + self.hc.hc_ffn_scale, + self.hc.hc_ffn_base, + ) + ffn_input = self.post_attention_layernorm(ffn_input) + ffn_output = self.mlp(ffn_input, input_ids) if isinstance(self.mlp, DeepseekV4MoE) else self.mlp(ffn_input) + hidden = self.hc.hc_post(ffn_output, residual_ffn, post_ffn, comb_ffn) + # Native C++ resets its optional residual at the start of every layer. + return hidden, None + + +class DeepseekV4Model(nn.Module): + """DeepSeek-V4 transformer body.""" + + def __init__(self, cfg: DeepseekV4Config, dtype: torch.dtype, device: torch.device) -> None: + super().__init__() + self.cfg = cfg + if cfg.cp_size > 1: + raise NotImplementedError("DeepSeek-V4 Python CP is reserved for the CP context PR") + tp = cfg.tp_size + self.embed_tokens = HiddenParallelEmbedding( + cfg.vocab_size, cfg.hidden_size // tp, tp, dtype=dtype, device=device + ) + self.layers = nn.ModuleList([DeepseekV4DecoderLayer(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) + # Model-level HyperConnection head: merges the hc_mult residual streams + # back into a single hidden vector before the final norm. Unlike the + # per-layer hc_pre (which uses mix_hc=(2+mult)*mult), the head uses a + # plain [hc_mult, hc_dim] hc_fn + [hc_mult] base + [1] scale. + hc_dim = cfg.hc_mult * cfg.hidden_size + self.hc_head_fn = nn.Parameter(torch.empty(cfg.hc_mult, hc_dim, dtype=torch.float32, device=device)) + self.hc_head_base = nn.Parameter(torch.empty(cfg.hc_mult, dtype=torch.float32, device=device)) + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32, device=device)) + # Native C++ falls back to max_position_embeddings when the flat + # rope_scaling_original_max_position_embeddings ModelArgs field is 0. + # The DSV4 loader currently leaves that field at 0, so old_context_len + # is 1048576 for this checkpoint, not nested rope_scaling's 65536. + native_old_context_len = cfg.max_position_embeddings + self.rotary = DeepseekV4RotaryEmbedding( + cfg.qk_rope_head_dim, + cfg.max_position_embeddings, + cfg.rope_scaling_factor, + cfg.rope_theta, + cfg.rope_beta_fast, + cfg.rope_beta_slow, + native_old_context_len, + dtype=dtype, + device=device, + ) + # Per-ratio compressed RoPE caches (C++ DeepseekV4RotaryEmbedding c4/c128 + # groups). Same YaRN inv_freq as the default cache but with + # compress_rope_theta (config=160000) and NO mscale amplitude (C++ + # create_cos_sin_cache does (void)mscale). mscale=1/mscale_all_dim=1 makes + # rope_mscale = get_mscale(s,1)/get_mscale(s,1) = 1.0 (no amplitude). + self.compress_rotary_c4 = DeepseekV4RotaryEmbedding( + cfg.qk_rope_head_dim, + cfg.max_position_embeddings, + cfg.rope_scaling_factor, + cfg.compress_rope_theta, + cfg.rope_beta_fast, + cfg.rope_beta_slow, + native_old_context_len, + dtype=dtype, + device=device, + ) + self.compress_rotary_c128 = DeepseekV4RotaryEmbedding( + cfg.qk_rope_head_dim, + cfg.max_position_embeddings, + cfg.rope_scaling_factor, + cfg.compress_rope_theta, + cfg.rope_beta_fast, + cfg.rope_beta_slow, + native_old_context_len, + dtype=dtype, + device=device, + ) + + def attach_rope_tables_to_backend( + self, + backend, + positions: torch.Tensor, + graph_bt_cols: int = 0, + metadata=None, + ) -> None: + """Attach default + per-ratio compressed RoPE caches to the backend. + + Called inside model forward after embedding, matching the C++ DSA + metadata construction order. Default cache uses rope_theta; c4/c128 use + compress_rope_theta with no mscale. + """ + if backend is None or not hasattr(backend, "attach_rope_tables"): + return + positions = positions.to(torch.int64).contiguous() + backend.attach_rope_tables( + positions, + self.rotary.cos_sin_cache, + graph_bt_cols=graph_bt_cols, + c4_cos_sin=self.compress_rotary_c4.cos_sin_cache, + c128_cos_sin=self.compress_rotary_c128.cos_sin_cache, + metadata=metadata, + ) + + def _hc_head(self, x: torch.Tensor) -> torch.Tensor: + """Final HyperConnection head. + + Matches C++ DeepseekV4ModelImpl::hc_head: this final merge is not the + per-layer Sinkhorn hc_pre kernel. Its checkpoint weights have shapes + hc_head_fn=[hc_mult, hc_mult*hidden], hc_head_base=[hc_mult], and + hc_head_scale=[1]. + """ + x_float = x.to(torch.float32) + x_flatten = x_float.flatten(-2, -1) + rsqrt = torch.rsqrt(x_flatten.pow(2).mean(-1, keepdim=True) + self.cfg.rms_norm_eps) + mixes = torch.matmul(x_flatten, self.hc_head_fn.transpose(0, 1)) + mixes = mixes * rsqrt + pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.cfg.hc_eps + y = (pre.unsqueeze(-1) * x_float).sum(-2) + return y.to(x.dtype) + + def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + hidden = self.embed_tokens(input_ids) + positions = positions.to(torch.int64).contiguous() + cos_sin_cache = self.rotary.cos_sin_cache + context = get_forward_context() + backend = context.attention_backend + metadata = context.metadata + self.attach_rope_tables_to_backend(backend, positions, metadata=metadata) + prepare_dsa = getattr(backend, "prepare_dsa_metadata_for_forward", None) + if prepare_dsa is not None: + prepare_dsa(metadata) + if self.cfg.cp_size > 1: + raise NotImplementedError("DeepSeek-V4 Python CP is reserved for the CP context PR") + # Expand hidden into hc_mult parallel residual streams for the + # HyperConnection decoder layers (C++ flat_hc does this reshape). + hidden = hidden.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1).contiguous() + residual: torch.Tensor | None = None + for layer_id, layer in enumerate(self.layers): + compress_ratio = self.cfg.compress_ratios[layer_id] if layer_id < len(self.cfg.compress_ratios) else 1 + if compress_ratio == 4: + layer_cos_sin_cache = self.compress_rotary_c4.cos_sin_cache + elif compress_ratio == 128: + layer_cos_sin_cache = self.compress_rotary_c128.cos_sin_cache + else: + layer_cos_sin_cache = cos_sin_cache + select_layer_rope = getattr(backend, "select_dsa_layer_rope", None) + if select_layer_rope is not None: + select_layer_rope(layer_id, layer_cos_sin_cache, metadata) + hidden, residual = layer( + hidden, + residual, + positions, + layer_cos_sin_cache, + input_ids, + ) + record_layer_event(layer_id) + # hc_head: merge the hc_mult streams back into a single hidden vector. + merged = self._hc_head(residual if residual is not None else hidden) + hidden = self.norm(merged, None) + return hidden + + +class DeepseekV4ForCausalLM(PyModelBase): + """DeepSeek-V4 causal LM driven by the C++ PyCausalLM bridge.""" + + def __init__(self, config: dict) -> None: + super().__init__() + self.cfg = DeepseekV4Config.from_dict(config) + if self.cfg.cp_size > 1: + raise NotImplementedError("DeepSeek-V4 Python CP is reserved for the CP context PR") + dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype")) + device = torch.device(config.get("device", "npu:0")) + self.model = DeepseekV4Model(self.cfg, dtype, device) + tp = self.cfg.tp_size + self.lm_head = ColumnParallelLinear( + self.cfg.hidden_size, + self.cfg.vocab_size // tp, + tp, + gather_output=True, + dtype=dtype, + device=device, + ) + + def load_weights(self, state_dicts, tp_rank: int, tp_size: int) -> None: + cfg = self.cfg + loader = W8A8WeightLoader(self, state_dicts, cfg.tp_size, cfg.tp_rank) + + def _has(name: str) -> bool: + return loader.find(name) is not None + + def _w8a8(ckpt_prefix: str, param_prefix: str, shard_dims: dict | None = None) -> None: + """Load a W8A8-dynamic projection. + + DSV4's checkpoint stores ``weight`` (int8) + ``weight_scale`` + + ``weight_offset`` (per-output-channel), which is the + ``W8A8DynamicLinear`` format -- NOT the static deq_scale/quant_bias + format of ``W8A8StaticLinear``. + """ + for suffix in ("weight", "weight_scale", "weight_offset"): + ckpt_key = ckpt_prefix + "." + suffix + if not _has(ckpt_key): + continue + t = loader.load_tensor(ckpt_key) + dim = (shard_dims or {}).get(suffix) + if dim is not None: + t = loader.shard(t, dim=dim) + loader.copy_in(param_prefix + "." + suffix, t) + + # --- Embedding (checkpoint: embed.weight). --- + loader.copy_in( + "model.embed_tokens.weight", + loader.shard(loader.load_tensor("embed.weight"), dim=1), + ) + + # --- Per-layer weights (checkpoint: layers.N.<...>). --- + for i in range(cfg.n_layers): + ck = f"layers.{i}." # checkpoint prefix + pm = f"model.layers.{i}." # parameter prefix + attn = self.model.layers[i].self_attn + # Attention W8A8 projections (ckpt name -> module name). + _w8a8(ck + "attn.wq_a", pm + "self_attn.q_a_proj") + _w8a8( + ck + "attn.wq_b", + pm + "self_attn.q_b_proj", + {"weight": 0, "weight_scale": 0, "weight_offset": 0}, + ) + _w8a8(ck + "attn.wkv", pm + "self_attn.kv_proj") + # o_a/o_b are bf16 (unquantized) column/row-parallel weights, not W8A8. + loader.copy_in( + pm + "self_attn.o_a_proj.weight", + loader.shard(loader.load_tensor(ck + "attn.wo_a.weight"), dim=0), + ) + loader.copy_in( + pm + "self_attn.o_b_proj.weight", + loader.shard(loader.load_tensor(ck + "attn.wo_b.weight"), dim=1), + ) + # Attention layernorms + sink. + loader.copy_in( + pm + "self_attn.q_a_layernorm.weight", + loader.load_tensor(ck + "attn.q_norm.weight"), + ) + loader.copy_in( + pm + "self_attn.kv_a_layernorm.weight", + loader.load_tensor(ck + "attn.kv_norm.weight"), + ) + # attn_sink (parameter): load either bare tensor or .weight form. + sink_key = ck + "attn.attn_sink" + if not _has(sink_key): + sink_key = ck + "attn.attn_sink.weight" + if _has(sink_key): + sink = loader.load_tensor(sink_key) + if sink.dim() == 1 and sink.size(0) == cfg.n_heads and cfg.tp_size > 1: + shard_size = cfg.n_heads // cfg.tp_size + sink = sink.narrow(0, cfg.tp_rank * shard_size, shard_size) + loader.copy_in(pm + "self_attn.attn_sink", sink) + # Layer layernorms (ckpt attn_norm/ffn_norm -> input/post_attention). + loader.copy_in( + pm + "input_layernorm.weight", + loader.load_tensor(ck + "attn_norm.weight"), + ) + loader.copy_in( + pm + "post_attention_layernorm.weight", + loader.load_tensor(ck + "ffn_norm.weight"), + ) + # HyperConnection weights (ckpt layers.N.hc_* -> model.layers.N.hc.hc_*). + for part in ("attn", "ffn"): + for suffix in ("fn", "scale", "base"): + name = f"hc_{part}_{suffix}" + loader.copy_in(pm + "hc." + name, loader.load_tensor(ck + name)) + # Indexer weights (ckpt layers.N.attn.indexer.*). + if attn.indexer is not None and _has(ck + "attn.indexer.wq_b.weight"): + # Indexer wq_b (ReplicatedLinear, not sharded) + weights_proj. + _w8a8(ck + "attn.indexer.wq_b", pm + "self_attn.indexer.wq_b") + loader.copy_in( + pm + "self_attn.indexer.weights_proj.weight", + loader.load_tensor(ck + "attn.indexer.weights_proj.weight"), + ) + # Compressor sub-module: wkv (unquantized f32 fused wk+wv) + + # wgate + ape + norm (all f32, not W8A8). + loader.copy_in( + pm + "self_attn.indexer.compressor_wkv.weight", + loader.load_tensor(ck + "attn.indexer.compressor.wkv.weight"), + ) + loader.copy_in( + pm + "self_attn.indexer.compressor_wgate.weight", + loader.load_tensor(ck + "attn.indexer.compressor.wgate.weight"), + ) + loader.copy_in( + pm + "self_attn.indexer.compressor_ape", + loader.load_tensor(ck + "attn.indexer.compressor.ape"), + ) + loader.copy_in( + pm + "self_attn.indexer.compressor_norm.weight", + loader.load_tensor(ck + "attn.indexer.compressor.norm.weight"), + ) + # Attention-level cmp_kv compressor (head_dim=512, separate from the + # indexer compressor at head_dim=128). Ckpt: attn.compressor.*. + # Mirrors C++ DSAttentionImpl compressor_ (compressor.cpp:590-597). + if hasattr(attn, "cmp_wkv") and _has(ck + "attn.compressor.wkv.weight"): + _w = loader.load_tensor(ck + "attn.compressor.wkv.weight") + loader.copy_in(pm + "self_attn.cmp_wkv.weight", _w) + loader.copy_in( + pm + "self_attn.cmp_wgate.weight", + loader.load_tensor(ck + "attn.compressor.wgate.weight"), + ) + loader.copy_in( + pm + "self_attn.cmp_ape", + loader.load_tensor(ck + "attn.compressor.ape"), + ) + loader.copy_in( + pm + "self_attn.cmp_norm.weight", + loader.load_tensor(ck + "attn.compressor.norm.weight"), + ) + attn.process_weights_after_loading() + # MoE / dense MLP weights -- DSV4 MoE uses hash routing (gate.weight + # + gate.tid2eid) and per-expert w1/w2/w3, distinct from DSV3.2's + # noaux_tc grouped_moe. Staged by _load_dsv4_moe below when the MoE + # adapter is in place; for now load the gate so the layer builds. + mlp = self.model.layers[i].mlp + if hasattr(mlp, "experts_w13") and _has(ck + "ffn.experts.0.w1.weight"): + self._load_dsv4_moe(loader, ck, pm, i) + mlp.process_weights_after_loading() + + # --- Final norm + hc_head + lm_head. --- + loader.copy_in("model.norm.weight", loader.load_tensor("norm.weight")) + loader.copy_in("model.hc_head_fn", loader.load_tensor("hc_head_fn")) + loader.copy_in("model.hc_head_base", loader.load_tensor("hc_head_base")) + loader.copy_in("model.hc_head_scale", loader.load_tensor("hc_head_scale")) + # Match LlmForCausalLMImplBase's non-tied output-head lookup order. The + # Flash checkpoint uses ``head.weight`` rather than ``lm_head.weight``. + lm_head_key = next( + ( + name + for name in ( + "lm_head.weight", + "model.lm_head.weight", + "model.head.weight", + "head.weight", + ) + if _has(name) + ), + None, + ) + assert lm_head_key is not None, "checkpoint output-head weight not found" + loader.copy_in( + "lm_head.weight", + loader.shard(loader.load_tensor(lm_head_key), dim=0), + ) + + def _load_dsv4_moe(self, loader, ck: str, pm: str, layer_id: int) -> None: + """Stage DSV4 MoE weights for DeepseekV4MoE (hash routing + EP sharding). + + Loads gate.weight + tid2eid (hash layers) + per-expert w1/w2/w3 + (only local EP experts, fused into w13=w1+w3) + shared_experts. + Mirrors C++ FusedMoEImpl::load_experts (fused_moe.cpp:1938+). + """ + + def _has(name: str) -> bool: + return loader.find(name) is not None + + cfg = self.cfg + mlp = self.model.layers[layer_id].mlp + # Gate weight [n_total_experts, hidden] float32 (replicated, not EP-sharded). + loader.copy_in(pm + "mlp.gate.weight", loader.load_tensor(ck + "ffn.gate.weight")) + if mlp.hash_layer: + # C++ DeepseekV4GateImpl requires tid2eid for every hash layer. + tid2eid_key = ck + "ffn.gate.tid2eid" + if not _has(tid2eid_key): + tid2eid_key += ".weight" + assert _has(tid2eid_key), f"hash gate checkpoint tensor not found: {tid2eid_key}" + loader.copy_in(pm + "mlp.tid2eid", loader.load_tensor(tid2eid_key)) + else: + # Match DeepseekV4GateImpl::load_state_dict: the correction bias is + # mandatory for non-hash routing, with the legacy key as fallback. + bias_key = ck + "ffn.gate.bias" + if not _has(bias_key): + bias_key = ck + "ffn.gate.e_score_correction_bias" + assert _has(bias_key), ( + f"non-hash gate checkpoint tensor not found: {ck}ffn.gate.bias (or e_score_correction_bias)" + ) + loader.copy_in( + pm + "mlp.e_score_correction_bias", + loader.load_tensor(bias_key), + ) + # Per-expert w1+w3 -> fused w13, w2 -> w2 (int8 + scale). + # EP: only load local experts [start_expert_id, start_expert_id + num_experts_per_rank). + tp = mlp.moe_tp_size + tp_rank = mlp.moe_tp_rank + start = mlp.start_expert_id + nepr = mlp.num_experts_per_rank + w13 = self.get_parameter(pm + "mlp.experts_w13") + w2 = self.get_parameter(pm + "mlp.experts_w2") + w13_scale = self.get_buffer(pm + "mlp.experts_w13_scale") + w2_scale = self.get_buffer(pm + "mlp.experts_w2_scale") + for local_idx in range(nepr): + global_id = start + local_idx + e = ck + f"ffn.experts.{global_id}." + w1 = loader.load_tensor(e + "w1.weight") + w3 = loader.load_tensor(e + "w3.weight") + w13_j = torch.cat([w1, w3], dim=0) + w2_j = loader.load_tensor(e + "w2.weight") + if tp > 1: + w13_j = loader.shard(w13_j, dim=0, world=tp, rank=tp_rank) + w2_j = loader.shard(w2_j, dim=1, world=tp, rank=tp_rank) + w13[local_idx].copy_(w13_j.to(w13.dtype)) + w2[local_idx].copy_(w2_j.to(w2.dtype)) + if _has(e + "w1.weight_scale"): + s1 = loader.load_tensor(e + "w1.weight_scale") + s3 = loader.load_tensor(e + "w3.weight_scale") if _has(e + "w3.weight_scale") else s1 + s13 = torch.cat([s1, s3], dim=0) + if tp > 1: + s13 = loader.shard(s13, dim=0, world=tp, rank=tp_rank) + w13_scale[local_idx].copy_(s13[: w13_j.size(0)]) + if _has(e + "w2.weight_scale"): + w2_scale[local_idx].copy_(loader.load_tensor(e + "w2.weight_scale")) + # Shared experts: checkpoint has w1/w2/w3 (W8A8 dynamic), fuse w1+w3 -> gate_up_proj. + se = ck + "ffn.shared_experts." + if _has(se + "w1.weight"): + se_w1 = loader.load_tensor(se + "w1.weight") + se_w3 = loader.load_tensor(se + "w3.weight") + se_w13 = torch.cat([se_w1, se_w3], dim=0) + if tp > 1: + se_w13 = loader.shard(se_w13, dim=0, world=tp, rank=tp_rank) + loader.copy_in(pm + "mlp.shared_experts.gate_up_proj.weight", se_w13) + if _has(se + "w1.weight_scale"): + s1 = loader.load_tensor(se + "w1.weight_scale") + s3 = loader.load_tensor(se + "w3.weight_scale") + se_s13 = torch.cat([s1, s3], dim=0) + if tp > 1: + se_s13 = loader.shard(se_s13, dim=0, world=tp, rank=tp_rank) + loader.copy_in(pm + "mlp.shared_experts.gate_up_proj.weight_scale", se_s13[: se_w13.size(0)]) + if _has(se + "w1.weight_offset"): + o1 = loader.load_tensor(se + "w1.weight_offset") + o3 = loader.load_tensor(se + "w3.weight_offset") + loader.copy_in( + pm + "mlp.shared_experts.gate_up_proj.weight_offset", torch.cat([o1, o3], dim=0)[: se_w13.size(0)] + ) + se_w2 = loader.load_tensor(se + "w2.weight") + if tp > 1: + se_w2 = loader.shard(se_w2, dim=1, world=tp, rank=tp_rank) + loader.copy_in(pm + "mlp.shared_experts.down_proj.weight", se_w2) + if _has(se + "w2.weight_scale"): + loader.copy_in( + pm + "mlp.shared_experts.down_proj.weight_scale", loader.load_tensor(se + "w2.weight_scale") + ) + if _has(se + "w2.weight_offset"): + loader.copy_in( + pm + "mlp.shared_experts.down_proj.weight_offset", loader.load_tensor(se + "w2.weight_offset") + ) + # NOTE: shared_experts.{gate_up,down}_proj.process_weights_after_loading + # is NOT called here. It is called exactly once via + # DeepseekV4MoE.process_weights_after_loading (line ~816) at the end of + # the per-layer load loop (load_weights line ~1187). Calling it here + # too would transpose the weight twice (process_weights is not + # idempotent), leaving it in [out, in] layout and tripping quant_matmul + # "x1 dim[-1] must match x2 dim[-2], got 4096 vs 512". diff --git a/xllm/python/registry.py b/xllm/python/registry.py index 2f8cb52090..d288a8b3e8 100644 --- a/xllm/python/registry.py +++ b/xllm/python/registry.py @@ -26,15 +26,15 @@ class by the model's architecture (or model_type) string. import torch.nn as nn _ModelPath = tuple[str, str] -_REGISTRY: Dict[str, _ModelPath] = {} +_REGISTRY: dict[str, _ModelPath] = {} def register_model( *names: str, -) -> Callable[[Type[nn.Module]], Type[nn.Module]]: +) -> Callable[[type[nn.Module]], type[nn.Module]]: """Register a model class for callers that already imported its module.""" - def deco(cls: Type[nn.Module]) -> Type[nn.Module]: + def deco(cls: type[nn.Module]) -> type[nn.Module]: path = (cls.__module__, cls.__name__) for name in names: _REGISTRY[name] = path @@ -49,11 +49,9 @@ def _register_model_path(module_name: str, class_name: str, *names: str) -> None _REGISTRY[name] = path -def get_model_class(name: str) -> Type[nn.Module]: +def get_model_class(name: str) -> type[nn.Module]: if name not in _REGISTRY: - raise KeyError( - f"model '{name}' not registered; available: {sorted(_REGISTRY)}" - ) + raise KeyError(f"model '{name}' not registered; available: {sorted(_REGISTRY)}") module_name, class_name = _REGISTRY[name] model_cls = getattr(import_module(module_name), class_name) return model_cls @@ -93,6 +91,12 @@ def _register_builtin_models() -> None: "Glm52ForCausalLM", "glm_moe_dsa", ) + _register_model_path( + "xllm.python.models.deepseek_v4", + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLM", + "deepseek_v4", + ) _register_model_path( "xllm.python.models.deepseek_v32_mtp",