From 78301a3a36494be20edc8ccf9b98b47ad9ee8ad0 Mon Sep 17 00:00:00 2001 From: Super User Date: Thu, 13 Aug 2026 15:57:02 +0800 Subject: [PATCH 1/3] feat: bridge DeepSeek V4 Python caches and metadata --- tests/python/test_model_executor.py | 49 ++++++++ xllm/core/runtime/py_attention_metadata.cpp | 128 ++++++++++++++++++++ xllm/core/runtime/py_attention_metadata.h | 25 ++++ xllm/core/runtime/py_executor_impl.cpp | 27 ++++- xllm/python/attention/backend.py | 32 ++++- 5 files changed, 255 insertions(+), 6 deletions(-) diff --git a/tests/python/test_model_executor.py b/tests/python/test_model_executor.py index 0f46882565..6d933c1785 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 @@ -604,6 +605,54 @@ 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/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..2691cc4426 100644 --- a/xllm/core/runtime/py_executor_impl.cpp +++ b/xllm/core/runtime/py_executor_impl.cpp @@ -45,6 +45,13 @@ 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; @@ -126,11 +133,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/python/attention/backend.py b/xllm/python/attention/backend.py index bf6ab94c47..96778ab013 100644 --- a/xllm/python/attention/backend.py +++ b/xllm/python/attention/backend.py @@ -42,10 +42,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, ...] @@ -84,6 +104,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] From 3ff984362f0af92d568566412225e7a713acfd8c Mon Sep 17 00:00:00 2001 From: Super User Date: Fri, 14 Aug 2026 01:07:21 +0800 Subject: [PATCH 2/3] feat: add DeepSeek V4 NPU kernel Python bindings --- tests/core/kernels/npu/npu_xllm_ops_test.cpp | 365 +++++++++++++++++ tests/python/test_kernels_import.py | 117 +++++- xllm/core/kernels/npu/npu_ops_library.cpp | 102 +++++ xllm/python/kernels_npu/__init__.py | 24 ++ xllm/python/kernels_npu/_custom_op.py | 400 +++++++++++++++++++ xllm/python/kernels_npu/dsa.py | 355 ++++++++++++++++ xllm/python/kernels_npu/normalization.py | 9 +- xllm/python/kernels_npu/rotary_embedding.py | 29 ++ 8 files changed, 1398 insertions(+), 3 deletions(-) create mode 100644 xllm/python/kernels_npu/dsa.py diff --git a/tests/core/kernels/npu/npu_xllm_ops_test.cpp b/tests/core/kernels/npu/npu_xllm_ops_test.cpp index 3a47ab700d..ff8a6d453f 100644 --- a/tests/core/kernels/npu/npu_xllm_ops_test.cpp +++ b/tests/core/kernels/npu/npu_xllm_ops_test.cpp @@ -209,6 +209,371 @@ TEST_F(NpuXllmOpsTest, EmbeddedInterpreterSeesOps) { .item(); } +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_kernels_import.py b/tests/python/test_kernels_import.py index 7fd8a98b5e..cd555fc445 100644 --- a/tests/python/test_kernels_import.py +++ b/tests/python/test_kernels_import.py @@ -79,6 +79,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( @@ -201,11 +212,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: @@ -237,6 +249,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/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index 1c9ae56305..b7ed5d5d4e 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, " @@ -349,11 +362,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)); @@ -379,4 +465,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/python/kernels_npu/__init__.py b/xllm/python/kernels_npu/__init__.py index dffb94fe5d..e931a38999 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, @@ -61,6 +72,7 @@ fused_add_rms_norm, l2_norm, rms_norm, + rms_norm_dynamic_quant, rms_norm_gated, ) from .quantization import ( @@ -73,6 +85,7 @@ interleaved_rotary_embedding, mrope, vision_rotary_mul, + npu_inplace_partial_rotary_mul, ) from .sparse_attention import ( lightning_indexer, @@ -85,6 +98,7 @@ __all__ = [ "rms_norm", "fused_add_rms_norm", + "rms_norm_dynamic_quant", "l2_norm", "rms_norm_gated", "silu_and_mul", @@ -93,6 +107,7 @@ "vision_fusion_attention", "fused_qk_norm_rope", "interleaved_rotary_embedding", + "npu_inplace_partial_rotary_mul", "mrope", "vision_rotary_mul", "moe_fused_topk", @@ -112,6 +127,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..a63e19b751 100644 --- a/xllm/python/kernels_npu/_custom_op.py +++ b/xllm/python/kernels_npu/_custom_op.py @@ -328,6 +328,334 @@ 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) @@ -345,3 +673,75 @@ def _sparse_flash_attention_out_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::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..0cf8db3a8d --- /dev/null +++ b/xllm/python/kernels_npu/dsa.py @@ -0,0 +1,355 @@ +# 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/normalization.py b/xllm/python/kernels_npu/normalization.py index e529abf6ef..e7843f2dc8 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: @@ -63,4 +64,10 @@ def rms_norm_gated( ) -__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..3d20346be7 100644 --- a/xllm/python/kernels_npu/rotary_embedding.py +++ b/xllm/python/kernels_npu/rotary_embedding.py @@ -175,9 +175,38 @@ def vision_rotary_mul( ).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__ = [ "fused_qk_norm_rope", "interleaved_rotary_embedding", "mrope", "vision_rotary_mul", + "npu_inplace_partial_rotary_mul", ] From fb1c6a1ebcc427fc1c5a03b42e108fb55e69e756 Mon Sep 17 00:00:00 2001 From: Super User Date: Fri, 14 Aug 2026 01:21:38 +0800 Subject: [PATCH 3/3] feat: add DeepSeek V4 DSA metadata builder --- tests/python/test_dsa_metadata.py | 434 ++++++++++++++++ xllm/python/attention/dsa_metadata.py | 720 ++++++++++++++++++++++++++ 2 files changed, 1154 insertions(+) create mode 100644 tests/python/test_dsa_metadata.py create mode 100644 xllm/python/attention/dsa_metadata.py diff --git a/tests/python/test_dsa_metadata.py b/tests/python/test_dsa_metadata.py new file mode 100644 index 0000000000..26a153e574 --- /dev/null +++ b/tests/python/test_dsa_metadata.py @@ -0,0 +1,434 @@ +# 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/xllm/python/attention/dsa_metadata.py b/xllm/python/attention/dsa_metadata.py new file mode 100644 index 0000000000..c8f46b6f6d --- /dev/null +++ b/xllm/python/attention/dsa_metadata.py @@ -0,0 +1,720 @@ +# 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 " + f"{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( + "ACL graph multi_block_tables must be 2-D: " + f"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