diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f0c68a14f..1090b90b5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,20 +3,7 @@ # It includes linting and platform-specific testing stages. name: CI -on: - schedule: - - cron: '0 18 * * *' # daily at 2:00 AM CST (UTC+8) - pull_request: - branches: [main, devops] - paths-ignore: - - "**.md" - - "docs/**" - - "examples/**" - - "docker/**" - - "LICENSE" - - ".github/ISSUE_TEMPLATE/**" - - ".github/PULL_REQUEST_TEMPLATE.md" - workflow_dispatch: +on: [] # disabled concurrency: group: ${{ github.workflow }}-${{ github.ref }} diff --git a/examples/run_dsv4.sh b/examples/run_dsv4.sh new file mode 100644 index 000000000..980aadb45 --- /dev/null +++ b/examples/run_dsv4.sh @@ -0,0 +1,14 @@ +export USE_FLAGGEMS=false +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 + +vllm serve /public-nvme/models/DeepSeek-V4-Flash-0731-INT \ + --trust-remote-code \ + --kv-cache-dtype fp8 \ + --block-size 256 \ + --enable-expert-parallel \ + --tensor-parallel-size 8 \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --reasoning-config '{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}' diff --git a/tests/unit_tests/compilation/test_deepseek_v4_splitting_ops.py b/tests/unit_tests/compilation/test_deepseek_v4_splitting_ops.py new file mode 100644 index 000000000..217613e71 --- /dev/null +++ b/tests/unit_tests/compilation/test_deepseek_v4_splitting_ops.py @@ -0,0 +1,22 @@ +# Copyright (c) 2025 BAAI. All rights reserved. + +from vllm.config.compilation import CompilationConfig, CompilationMode + +import vllm_fl + + +def test_register_deepseek_v4_fl_attention_as_splitting_op(monkeypatch): + attention_ops = ["vllm::deepseek_v4_attention"] + monkeypatch.setattr(CompilationConfig, "_attention_ops", attention_ops) + + vllm_fl._register_compilation_splitting_ops() + vllm_fl._register_compilation_splitting_ops() + + assert attention_ops == [ + "vllm::deepseek_v4_attention", + "vllm::deepseek_v4_fl_attention", + ] + + config = CompilationConfig(mode=CompilationMode.VLLM_COMPILE) + config.set_splitting_ops_for_v1("") + assert config.splitting_ops.count("vllm::deepseek_v4_fl_attention") == 1 diff --git a/tests/unit_tests/dispatch/test_cuda_backend_detection.py b/tests/unit_tests/dispatch/test_cuda_backend_detection.py new file mode 100644 index 000000000..83baf58dd --- /dev/null +++ b/tests/unit_tests/dispatch/test_cuda_backend_detection.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""Tests for NVIDIA vendor backend detection on vLLM CUDA platforms.""" + +from unittest.mock import Mock + +import torch +from vllm import platforms + +from vllm_fl.dispatch.backends.vendor.cuda.cuda import CudaBackend +from vllm_fl.dispatch.builtin_ops import _get_current_vendor_backend_dirs + + +def test_in_tree_cuda_platform_selects_cuda_vendor_backend(monkeypatch): + platform = Mock() + platform.vendor_name = None + platform.device_name = "cuda" + platform.is_cuda.return_value = True + monkeypatch.setattr(platforms, "current_platform", platform) + + assert _get_current_vendor_backend_dirs({"cuda", "ascend"}) == "cuda" + + +def test_cuda_backend_available_for_in_tree_cuda_platform(monkeypatch): + platform = Mock() + platform.device_name = "cuda" + platform.is_cuda.return_value = True + monkeypatch.setattr(platforms, "current_platform", platform) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 8) + monkeypatch.setattr(CudaBackend, "_available", None) + + assert CudaBackend().is_available() + + +def test_cuda_alike_platform_does_not_select_nvidia_backend(monkeypatch): + platform = Mock() + platform.vendor_name = None + platform.device_name = "cuda" + platform.is_cuda.return_value = False + monkeypatch.setattr(platforms, "current_platform", platform) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 8) + monkeypatch.setattr(CudaBackend, "_available", None) + + assert _get_current_vendor_backend_dirs({"cuda", "ascend"}) is None + assert not CudaBackend().is_available() diff --git a/tests/unit_tests/dispatch/test_deepseek_v4_attention.py b/tests/unit_tests/dispatch/test_deepseek_v4_attention.py new file mode 100644 index 000000000..7ae1cb5dd --- /dev/null +++ b/tests/unit_tests/dispatch/test_deepseek_v4_attention.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""Tests for the DeepSeek-V4 attention compile boundary.""" + +from types import SimpleNamespace + +import torch + +from vllm_fl.models import deepseek_v4 + + +def test_deepseek_v4_fl_attention_writes_preallocated_output(monkeypatch): + calls = [] + + class Layer: + def attention_impl(self, *args): + calls.append(args) + args[-1].fill_(7) + + layer = Layer() + monkeypatch.setattr( + deepseek_v4, + "get_forward_context", + lambda: SimpleNamespace(no_compile_layers={"layer": layer}), + ) + + tensors = [torch.empty(1) for _ in range(7)] + out = torch.empty(2, 3, 4) + result = deepseek_v4._deepseek_v4_fl_attention(*tensors, out, "layer") + + assert result is None + assert len(calls) == 1 + assert calls[0] == (*tensors, out) + assert calls[0][-1] is out + assert torch.equal(out, torch.full_like(out, 7)) + schema = torch._C._dispatch_find_schema_or_throw( + "vllm::deepseek_v4_fl_attention", "" + ).schema() + assert "Tensor(a7!) out" in str(schema) diff --git a/tests/unit_tests/dispatch/test_deepseek_v4_ops.py b/tests/unit_tests/dispatch/test_deepseek_v4_ops.py new file mode 100644 index 000000000..a2bfc8349 --- /dev/null +++ b/tests/unit_tests/dispatch/test_deepseek_v4_ops.py @@ -0,0 +1,218 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""Tests for DeepSeek-V4 operator dispatch.""" + +from unittest.mock import Mock + +import torch + +from vllm_fl.dispatch.backends.reference.impl.deepseek_v4 import ( + deepseek_v4_hc_head_torch, + deepseek_v4_int8_scaled_mm_torch, + deepseek_v4_inv_rope_quant_int8_torch, + deepseek_v4_mhc_post_torch, +) +from vllm_fl.dispatch.types import BackendImplKind +from vllm_fl.ops import deepseek_v4_int8_woa + +DSV4_OPS = { + "deepseek_v4_inv_rope_quant_int8", + "deepseek_v4_inv_rope_quant_fp8", + "deepseek_v4_int8_scaled_mm", + "deepseek_v4_mhc_pre", + "deepseek_v4_mhc_fused_post_pre", + "deepseek_v4_mhc_post", + "deepseek_v4_hc_head", + "deepseek_v4_fused_q_kv_rmsnorm", + "deepseek_v4_qnorm_rope_kv_quant_insert", + "deepseek_v4_qnorm_rope_kv_bf16_insert", + "deepseek_v4_qnorm_rope_kv_fp8_insert", + "deepseek_v4_compute_global_topk_indices_and_lens", + "deepseek_v4_flash_mla_with_kvcache", + "deepseek_v4_dequantize_and_gather_k_cache", + "deepseek_v4_combine_topk_swa_indices", + "deepseek_v4_flash_mla_sparse_fwd", + "deepseek_v4_fused_indexer_q_rope_quant", + "deepseek_v4_fused_indexer_q_rope_quant_int8", + "deepseek_v4_compress_int8_indexer_k_cache", + "deepseek_v4_int8_mqa_logits", + "deepseek_v4_int8_paged_mqa_logits", +} + +SPARSE_INDEXER_OPS = { + "indexer_k_quant_and_cache", + "cp_gather_indexer_k_quant_cache", + "top_k_per_row_prefill", + "top_k_per_row_decode", + "pack_seq_triton", + "unpack_seq_triton", +} + + +def test_reference_inv_rope_quant_int8(): + o = torch.tensor( + [[[1, 2, 3, 4], [-1, -2, 5, 6]]], + dtype=torch.bfloat16, + ) + positions = torch.tensor([0], dtype=torch.int32) + cos_sin_cache = torch.tensor([[0, 1]], dtype=torch.float32) + + quantized, scales = deepseek_v4_inv_rope_quant_int8_torch( + o, + positions, + cos_sin_cache, + n_groups=1, + heads_per_group=2, + nope_dim=2, + rope_dim=2, + ) + + expected = torch.tensor( + [[[21, 42, 85, -64, -21, -42, 127, -106]]], + dtype=torch.int8, + ) + assert torch.equal(quantized, expected) + torch.testing.assert_close( + scales, + torch.tensor([[[6 / 127]]], dtype=torch.float32), + ) + + +def test_frontend_dispatches_through_cached_op(monkeypatch): + expected = (Mock(), Mock()) + dispatch = Mock(return_value=expected) + monkeypatch.setattr( + deepseek_v4_int8_woa, + "_dispatch_inv_rope_quant_int8", + dispatch, + ) + args = ( + Mock(), + Mock(), + Mock(), + 2, + 4, + 64, + 64, + ) + + actual = deepseek_v4_int8_woa.fused_inv_rope_quant_int8(*args) + + assert actual is expected + dispatch.assert_called_once_with(*args) + + +def test_all_backends_register_deepseek_v4_op(monkeypatch): + from vllm_fl.dispatch.backends.flaggems import register_ops as flaggems_ops + from vllm_fl.dispatch.backends.reference import register_ops as reference_ops + from vllm_fl.dispatch.backends.vendor.cuda import register_ops as cuda_ops + + registered = [] + + class Registry: + def register_many(self, impls): + registered.extend(impls) + + monkeypatch.setattr( + flaggems_ops, + "use_flaggems_op", + lambda op_name: op_name == deepseek_v4_int8_woa.DSV4_INV_ROPE_QUANT_INT8_OP, + ) + registry = Registry() + flaggems_ops.register_builtins(registry) + cuda_ops.register_builtins(registry) + reference_ops.register_builtins(registry) + + implementations = [ + impl + for impl in registered + if impl.op_name == deepseek_v4_int8_woa.DSV4_INV_ROPE_QUANT_INT8_OP + ] + assert {impl.impl_id for impl in implementations} == { + "default.flagos", + "vendor.cuda", + "reference.torch", + } + assert {impl.kind for impl in implementations} == { + BackendImplKind.DEFAULT, + BackendImplKind.VENDOR, + BackendImplKind.REFERENCE, + } + + +def test_reference_scaled_mm_and_mhc_ops(): + x_q = torch.tensor([[1, -2]], dtype=torch.int8) + weight = torch.tensor([[3, 4], [5, 6]], dtype=torch.int8) + actual = deepseek_v4_int8_scaled_mm_torch( + x_q, + weight, + torch.tensor([[0.5]]), + torch.tensor([0.25, 0.5]), + torch.float32, + ) + torch.testing.assert_close(actual, torch.tensor([[-0.875, -2.0]])) + + residual = torch.tensor([[[1, 2], [3, 4]]], dtype=torch.bfloat16) + layer = torch.tensor([[2, -1]], dtype=torch.bfloat16) + post = torch.tensor([[[0.5], [1.0]]], dtype=torch.float32) + comb = torch.eye(2, dtype=torch.float32).unsqueeze(0) + torch.testing.assert_close( + deepseek_v4_mhc_post_torch(layer, residual, post, comb), + torch.tensor([[[2, 1.5], [5, 3]]], dtype=torch.bfloat16), + ) + + fn = torch.zeros((2, 4), dtype=torch.float32) + head = deepseek_v4_hc_head_torch( + residual, + fn, + torch.ones(1), + torch.zeros(2), + 1e-6, + 0.0, + ) + torch.testing.assert_close(head, residual.float().mean(dim=1).to(torch.bfloat16)) + + +def test_all_backends_register_all_deepseek_v4_ops(monkeypatch): + from vllm_fl.dispatch.backends.flaggems import register_ops as flaggems_ops + from vllm_fl.dispatch.backends.reference import register_ops as reference_ops + from vllm_fl.dispatch.backends.vendor.cuda import register_ops as cuda_ops + + registered = [] + + class Registry: + def register_many(self, impls): + registered.extend(impls) + + monkeypatch.setattr( + flaggems_ops, + "use_flaggems_op", + lambda op_name: op_name in DSV4_OPS | SPARSE_INDEXER_OPS, + ) + registry = Registry() + flaggems_ops.register_builtins(registry) + cuda_ops.register_builtins(registry) + reference_ops.register_builtins(registry) + + for op_name in DSV4_OPS: + implementations = [impl for impl in registered if impl.op_name == op_name] + assert {impl.impl_id for impl in implementations} == { + "default.flagos", + "vendor.cuda", + "reference.torch", + } + + for op_name in SPARSE_INDEXER_OPS: + implementations = [impl for impl in registered if impl.op_name == op_name] + assert {impl.impl_id for impl in implementations} == { + "default.flagos", + "vendor.cuda", + "reference.torch", + } + + +def test_sparse_indexer_overrides_upstream_cuda_entrypoint(): + from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer + from vllm_fl.ops.sparse_attn_indexer import SparseAttnIndexerFL + + assert SparseAttnIndexerFL.forward_cuda is not SparseAttnIndexer.forward_cuda diff --git a/tests/unit_tests/ops/test_deepseek_v4_int8_indexer.py b/tests/unit_tests/ops/test_deepseek_v4_int8_indexer.py new file mode 100644 index 000000000..ee7625aac --- /dev/null +++ b/tests/unit_tests/ops/test_deepseek_v4_int8_indexer.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""CUDA correctness tests for the DeepSeek-V4 INT8 indexer kernels.""" + +import pytest +import torch + +from vllm_fl.ops.deepseek_v4_int8_indexer import ( + int8_mqa_logits, + int8_paged_mqa_logits, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_int8_mqa_logits_matches_torch(): + torch.manual_seed(2) + num_queries, num_keys, num_heads, head_dim = 1, 97, 64, 128 + q = torch.randint( + -127, + 128, + (num_queries, num_heads, head_dim), + dtype=torch.int8, + ) + k = torch.randint(-127, 128, (num_keys, head_dim), dtype=torch.int8) + k_scale = torch.rand(num_keys, dtype=torch.float32) * 0.02 + weights = torch.randn(num_queries, num_heads, dtype=torch.float32) + cu_ks = torch.tensor([0], dtype=torch.int32) + cu_ke = torch.tensor([num_keys], dtype=torch.int32) + + actual = int8_mqa_logits( + q.cuda(), + k.cuda(), + k_scale.cuda(), + weights.cuda(), + cu_ks.cuda(), + cu_ke.cuda(), + ).cpu() + dots = torch.einsum("mhd,nd->mhn", q.float(), k.float()) + expected = (dots * k_scale[None, None, :]).relu() + expected = (expected * weights[:, :, None]).sum(dim=1) + + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-3) + + +def test_int8_paged_mqa_logits_matches_torch(): + torch.manual_seed(3) + batch, next_n, num_heads, head_dim = 1, 1, 64, 128 + block_size, num_blocks, context_len = 64, 2, 100 + q = torch.randint( + -127, + 128, + (batch, next_n, num_heads, head_dim), + dtype=torch.int8, + ) + k = torch.randint( + -127, + 128, + (num_blocks, block_size, head_dim), + dtype=torch.int8, + ) + k_scale = torch.rand(num_blocks, block_size, dtype=torch.float32) * 0.02 + weights = torch.randn(batch * next_n, num_heads, dtype=torch.float32) + + # The compressor stores one packed page as all INT8 K bytes followed by + # all fp32 scales. The logical tensor shape only reserves 132 bytes/token; + # its final dimension must not be interpreted as an interleaved layout. + cache = torch.empty( + num_blocks, + block_size, + head_dim + torch.tensor([], dtype=torch.float32).element_size(), + dtype=torch.uint8, + ) + flat_cache = cache.view(-1) + for block in range(num_blocks): + page_base = block * cache.stride(0) + k_bytes = k[block].contiguous().view(torch.uint8).reshape(-1) + scale_bytes = k_scale[block].contiguous().view(torch.uint8).reshape(-1) + flat_cache[page_base : page_base + k_bytes.numel()].copy_(k_bytes) + scale_start = page_base + k_bytes.numel() + flat_cache[scale_start : scale_start + scale_bytes.numel()].copy_(scale_bytes) + + context_lens = torch.tensor([[context_len]], dtype=torch.int32) + block_table = torch.tensor([[0, 1]], dtype=torch.int32) + actual = int8_paged_mqa_logits( + q.cuda(), + cache.cuda(), + weights.cuda(), + context_lens.cuda(), + block_table.cuda(), + num_blocks * block_size, + ).cpu() + + flat_k = k.reshape(-1, head_dim)[:context_len] + flat_scale = k_scale.reshape(-1)[:context_len] + dots = torch.einsum("hd,nd->hn", q[0, 0].float(), flat_k.float()) + expected = (dots * flat_scale[None, :]).relu() + expected = (expected * weights[0, :, None]).sum(dim=0) + + torch.testing.assert_close(actual[0, :context_len], expected, atol=2e-3, rtol=2e-3) + assert torch.isfinite(actual[0, :context_len]).all() diff --git a/tests/unit_tests/quantization/test_w8a8_linear.py b/tests/unit_tests/quantization/test_w8a8_linear.py index f5343632b..d7d7054d9 100644 --- a/tests/unit_tests/quantization/test_w8a8_linear.py +++ b/tests/unit_tests/quantization/test_w8a8_linear.py @@ -254,3 +254,51 @@ def fake_scaled_mm( assert actual.shape == (1, 2, 3) assert torch.equal(actual, expected) + + +def test_w8a8_grouped_linear_prepares_group_major_weights(): + checkpoint_weight = torch.arange(32, dtype=torch.int8).reshape(8, 4) + checkpoint_scale = torch.arange(1, 9, dtype=torch.float32).reshape(8, 1) + layer = torch.nn.Module() + layer.is_bmm = True + layer.bmm_batch_size = 2 + layer.register_parameter( + "weight", + torch.nn.Parameter(checkpoint_weight.clone(), requires_grad=False), + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter(checkpoint_scale.clone(), requires_grad=False), + ) + layer.register_parameter("input_scale", None) + layer.register_parameter("input_zero_point", None) + layer.register_parameter("azp_adj", None) + + kernel = object.__new__(linear.FLW8A8DynamicLinearKernel) + kernel.layer_param_names = [ + "weight", + "weight_scale", + "input_scale", + "input_zero_point", + "azp_adj", + ] + kernel.process_weights_after_loading(layer) + + grouped_weight = layer._fl_w8a8_grouped_weight + assert grouped_weight.shape == (2, 4, 4) + assert grouped_weight.is_contiguous() + assert torch.equal( + grouped_weight[0], + checkpoint_weight[:4].contiguous(), + ) + assert torch.equal( + grouped_weight[1], + checkpoint_weight[4:].contiguous(), + ) + assert grouped_weight[0].transpose(0, 1).stride() == (1, 4) + assert torch.equal( + layer._fl_w8a8_grouped_weight_scale, + checkpoint_scale.reshape(2, 4), + ) + assert "_fl_w8a8_grouped_weight" not in layer.state_dict() + assert "_fl_w8a8_grouped_weight_scale" not in layer.state_dict() diff --git a/vllm_fl/__init__.py b/vllm_fl/__init__.py index 42ade8165..716987249 100644 --- a/vllm_fl/__init__.py +++ b/vllm_fl/__init__.py @@ -25,6 +25,8 @@ logger = logging.getLogger(__name__) +_DEEPSEEK_V4_FL_ATTENTION_OP = "vllm::deepseek_v4_fl_attention" + def __getattr__(name): if name == "distributed": @@ -95,12 +97,26 @@ def _patch_custom_ops(): register_op_schemas() +def _register_compilation_splitting_ops(): + """Keep stateful FL attention outside compiled graph partitions.""" + from vllm.config.compilation import CompilationConfig + + if _DEEPSEEK_V4_FL_ATTENTION_OP not in CompilationConfig._attention_ops: + CompilationConfig._attention_ops.append(_DEEPSEEK_V4_FL_ATTENTION_OP) + + def register(): """Register the FL platform.""" _patch_custom_ops() _patch_flash_attn_import() _patch_transformers_compat() + # Platform plugins are evaluated while vllm.platforms is still resolving + # current_platform. Importing CompilationConfig here re-enters platform + # resolution on vLLM 0.24; the plugin loader suppresses that exception and + # silently falls back to the in-tree CUDA platform. Register model-specific + # splitting ops later from register_model(), after PlatformFL is active. + # Model-specific platform patches from vllm_fl.patches.glm_moe_dsa import apply_platform_patches as glm5_platform glm5_platform() @@ -140,6 +156,8 @@ def register_model(): """Register FL-specific models not yet upstream.""" # General plugins are loaded independently in spawned model-inspection and # worker processes, so all runtime compatibility hooks must be idempotent. + _register_compilation_splitting_ops() + from vllm_fl.patches.moe_sum import patch_vllm_moe_sum from vllm_fl.patches.qwen3_5_text import apply_qwen3_5_text_patches @@ -152,6 +170,15 @@ def register_model(): register_quant_linear() register_router() + # Replace only the DSV4 registry entry. The lazy thin model subclasses the + # vLLM 0.24 NVIDIA implementation and preserves its non-INT8 behavior. + from vllm import ModelRegistry + + ModelRegistry.register_model( + "DeepseekV4ForCausalLM", + "vllm_fl.models.deepseek_v4:DeepseekV4FLForCausalLM", + ) + # Register GLM-5 (GlmMoeDsa) — config not yet upstream try: from vllm.transformers_utils.config import _CONFIG_REGISTRY diff --git a/vllm_fl/dispatch/README.md b/vllm_fl/dispatch/README.md index 674301a18..87d20de65 100644 --- a/vllm_fl/dispatch/README.md +++ b/vllm_fl/dispatch/README.md @@ -525,6 +525,7 @@ Currently supported operators: | Operator | Description | FlagGems | Reference | Vendor | |----------|-------------|----------|-----------|--------| | `dynamic_per_token_quant_int8` | vLLM-compatible symmetric dynamic per-token INT8 quantization | ✓ | ✓ | - | +| `deepseek_v4_inv_rope_quant_int8` | DSV4 inverse-RoPE with group-major INT8 activation quantization | ✓ | ✓ | ✓ | | `silu_and_mul` | SiLU activation + element-wise multiplication | ✓ | ✓ | ✓ | | `rms_norm` | RMS normalization | ✓ | ✓ | ✓ | | `rotary_embedding` | Rotary position embedding | ✓ | ✓ | ✓ | diff --git a/vllm_fl/dispatch/backends/flaggems/flaggems.py b/vllm_fl/dispatch/backends/flaggems/flaggems.py index bc6077a91..a09f8ee47 100644 --- a/vllm_fl/dispatch/backends/flaggems/flaggems.py +++ b/vllm_fl/dispatch/backends/flaggems/flaggems.py @@ -8,10 +8,9 @@ from __future__ import annotations -from typing import Optional, Union +import os import torch -import os from vllm_fl.dispatch.backends.base import Backend @@ -24,7 +23,7 @@ class FlagGemsBackend(Backend): operator implementations. """ - _available: Optional[bool] = None + _available: bool | None = None @property def name(self) -> str: @@ -42,6 +41,127 @@ def is_available(self) -> bool: return FlagGemsBackend._available # ==================== Operator Implementations ==================== + def deepseek_v4_inv_rope_quant_int8( + self, + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .impl.deepseek_v4 import ( + deepseek_v4_inv_rope_quant_int8_flaggems, + ) + + return deepseek_v4_inv_rope_quant_int8_flaggems( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + + def _deepseek_v4_call(self, op_name, *args, **kwargs): + from .impl import deepseek_v4 + + fn = getattr(deepseek_v4, f"deepseek_v4_{op_name}_flaggems") + return fn(*args, **kwargs) + + def deepseek_v4_inv_rope_quant_fp8(self, *args, **kwargs): + return self._deepseek_v4_call("inv_rope_quant_fp8", *args, **kwargs) + + def deepseek_v4_int8_scaled_mm(self, *args, **kwargs): + return self._deepseek_v4_call("int8_scaled_mm", *args, **kwargs) + + def deepseek_v4_mhc_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_pre", *args, **kwargs) + + def deepseek_v4_mhc_fused_post_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_fused_post_pre", *args, **kwargs) + + def deepseek_v4_mhc_post(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_post", *args, **kwargs) + + def deepseek_v4_hc_head(self, *args, **kwargs): + return self._deepseek_v4_call("hc_head", *args, **kwargs) + + def deepseek_v4_fused_q_kv_rmsnorm(self, *args, **kwargs): + return self._deepseek_v4_call("fused_q_kv_rmsnorm", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_quant_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_quant_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_bf16_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_bf16_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_fp8_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_fp8_insert", *args, **kwargs) + + def deepseek_v4_compute_global_topk_indices_and_lens(self, *args, **kwargs): + return self._deepseek_v4_call( + "compute_global_topk_indices_and_lens", *args, **kwargs + ) + + def deepseek_v4_flash_mla_with_kvcache(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_with_kvcache", *args, **kwargs) + + def deepseek_v4_dequantize_and_gather_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call("dequantize_and_gather_k_cache", *args, **kwargs) + + def deepseek_v4_combine_topk_swa_indices(self, *args, **kwargs): + return self._deepseek_v4_call("combine_topk_swa_indices", *args, **kwargs) + + def deepseek_v4_flash_mla_sparse_fwd(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_sparse_fwd", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant(self, *args, **kwargs): + return self._deepseek_v4_call("fused_indexer_q_rope_quant", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant_int8(self, *args, **kwargs): + return self._deepseek_v4_call( + "fused_indexer_q_rope_quant_int8", *args, **kwargs + ) + + def deepseek_v4_compress_int8_indexer_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call( + "compress_int8_indexer_k_cache", *args, **kwargs + ) + + def deepseek_v4_int8_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_mqa_logits", *args, **kwargs) + + def deepseek_v4_int8_paged_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_paged_mqa_logits", *args, **kwargs) + + def _sparse_indexer_call(self, op_name, *args, **kwargs): + from .impl import sparse_attn_indexer + + fn = getattr(sparse_attn_indexer, f"{op_name}_flaggems") + return fn(*args, **kwargs) + + def indexer_k_quant_and_cache(self, *args, **kwargs): + return self._sparse_indexer_call("indexer_k_quant_and_cache", *args, **kwargs) + + def cp_gather_indexer_k_quant_cache(self, *args, **kwargs): + return self._sparse_indexer_call( + "cp_gather_indexer_k_quant_cache", *args, **kwargs + ) + + def top_k_per_row_prefill(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_prefill", *args, **kwargs) + + def top_k_per_row_decode(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_decode", *args, **kwargs) + + def pack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("pack_seq_triton", *args, **kwargs) + + def unpack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("unpack_seq_triton", *args, **kwargs) def dynamic_per_token_quant_int8( self, @@ -93,8 +213,8 @@ def rms_norm( self, obj, x: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ RMS normalization. @@ -182,7 +302,9 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) -> if use_flaggems_attn: print("Using FlagGems attention backend.") - return "vllm_fl.dispatch.backends.flaggems.impl.attention.AttentionFLBackend" + return ( + "vllm_fl.dispatch.backends.flaggems.impl.attention.AttentionFLBackend" + ) return AttentionBackendEnum.TRITON_ATTN.get_path() @@ -191,7 +313,7 @@ def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int, - expert_map: Optional[torch.Tensor] = None, + expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ignore_invalid_experts: bool = False, ): @@ -287,6 +409,12 @@ def grouped_topk( from .impl.fused_moe import grouped_topk_flaggems return grouped_topk_flaggems( - scores, n_group, topk_group, topk, - renormalize, routed_scaling_factor, bias, scoring_func, + scores, + n_group, + topk_group, + topk, + renormalize, + routed_scaling_factor, + bias, + scoring_func, ) diff --git a/vllm_fl/dispatch/backends/flaggems/impl/deepseek_v4.py b/vllm_fl/dispatch/backends/flaggems/impl/deepseek_v4.py new file mode 100644 index 000000000..687c97b27 --- /dev/null +++ b/vllm_fl/dispatch/backends/flaggems/impl/deepseek_v4.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""FlagGems implementation of DeepSeek-V4-specific operators.""" + +from __future__ import annotations + +import torch +from flag_gems.runtime import torch_device_fn + +from vllm_fl.ops.deepseek_v4_int8_woa import ( + fused_inv_rope_quant_int8_triton, +) + + +def deepseek_v4_inv_rope_quant_int8_flaggems( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the fused backend-neutral Triton kernel under FlagGems' device guard.""" + with torch_device_fn.device(o.device): + return fused_inv_rope_quant_int8_triton( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + + +def deepseek_v4_inv_rope_quant_fp8_flaggems(*args): + from vllm_fl.dispatch.backends.vendor.cuda.impl.deepseek_v4 import ( + deepseek_v4_inv_rope_quant_fp8_cuda, + ) + + with torch_device_fn.device(args[0].device): + return deepseek_v4_inv_rope_quant_fp8_cuda(*args) + + +def deepseek_v4_int8_scaled_mm_flaggems( + x, weight, scale_a, scale_b, out_dtype, bias=None +): + from flag_gems import scaled_mm + + with torch_device_fn.device(x.device): + return scaled_mm(x, weight, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + + +def _reference(name, *args): + from vllm_fl.dispatch.backends.reference.impl import deepseek_v4 + + with torch_device_fn.device(args[0].device): + return getattr(deepseek_v4, name)(*args) + + +def deepseek_v4_mhc_pre_flaggems(*args): + return _reference("deepseek_v4_mhc_pre_torch", *args) + + +def deepseek_v4_mhc_fused_post_pre_flaggems(*args): + return _reference("deepseek_v4_mhc_fused_post_pre_torch", *args) + + +def deepseek_v4_mhc_post_flaggems(*args): + return _reference("deepseek_v4_mhc_post_torch", *args) + + +def deepseek_v4_hc_head_flaggems(*args): + return _reference("deepseek_v4_hc_head_torch", *args) + + +def _reference_kwargs(name, *args, **kwargs): + from vllm_fl.dispatch.backends.reference.impl import deepseek_v4 + + tensor = next((arg for arg in args if isinstance(arg, torch.Tensor)), None) + if tensor is None: + tensor = next( + (value for value in kwargs.values() if isinstance(value, torch.Tensor)), + None, + ) + if tensor is None: + return getattr(deepseek_v4, name)(*args, **kwargs) + with torch_device_fn.device(tensor.device): + return getattr(deepseek_v4, name)(*args, **kwargs) + + +def deepseek_v4_fused_q_kv_rmsnorm_flaggems(*args, **kwargs): + return _reference_kwargs("deepseek_v4_fused_q_kv_rmsnorm_torch", *args, **kwargs) + + +def deepseek_v4_qnorm_rope_kv_quant_insert_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_qnorm_rope_kv_quant_insert_torch", *args, **kwargs + ) + + +def deepseek_v4_qnorm_rope_kv_bf16_insert_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_qnorm_rope_kv_bf16_insert_torch", *args, **kwargs + ) + + +def deepseek_v4_qnorm_rope_kv_fp8_insert_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_qnorm_rope_kv_fp8_insert_torch", *args, **kwargs + ) + + +def deepseek_v4_compute_global_topk_indices_and_lens_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_compute_global_topk_indices_and_lens_torch", *args, **kwargs + ) + + +def deepseek_v4_flash_mla_with_kvcache_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_flash_mla_with_kvcache_torch", *args, **kwargs + ) + + +def deepseek_v4_dequantize_and_gather_k_cache_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_dequantize_and_gather_k_cache_torch", *args, **kwargs + ) + + +def deepseek_v4_combine_topk_swa_indices_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_combine_topk_swa_indices_torch", *args, **kwargs + ) + + +def deepseek_v4_flash_mla_sparse_fwd_flaggems(*args, **kwargs): + return _reference_kwargs("deepseek_v4_flash_mla_sparse_fwd_torch", *args, **kwargs) + + +def deepseek_v4_fused_indexer_q_rope_quant_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_fused_indexer_q_rope_quant_torch", *args, **kwargs + ) + + +def deepseek_v4_fused_indexer_q_rope_quant_int8_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_fused_indexer_q_rope_quant_int8_torch", *args, **kwargs + ) + + +def deepseek_v4_compress_int8_indexer_k_cache_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_compress_int8_indexer_k_cache_torch", *args, **kwargs + ) + + +def deepseek_v4_int8_mqa_logits_flaggems(*args, **kwargs): + return _reference_kwargs("deepseek_v4_int8_mqa_logits_torch", *args, **kwargs) + + +def deepseek_v4_int8_paged_mqa_logits_flaggems(*args, **kwargs): + return _reference_kwargs( + "deepseek_v4_int8_paged_mqa_logits_torch", *args, **kwargs + ) diff --git a/vllm_fl/dispatch/backends/flaggems/impl/sparse_attn_indexer.py b/vllm_fl/dispatch/backends/flaggems/impl/sparse_attn_indexer.py new file mode 100644 index 000000000..f12a4b806 --- /dev/null +++ b/vllm_fl/dispatch/backends/flaggems/impl/sparse_attn_indexer.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""FlagGems wrappers for sparse indexer helper ops.""" + +from __future__ import annotations + +import torch +from flag_gems.runtime import torch_device_fn + + +def _native(op_name, *args, **kwargs): + from vllm_fl.dispatch.backends.reference.impl import sparse_attn_indexer + + fn = getattr(sparse_attn_indexer, f"{op_name}_torch") + tensor = next((arg for arg in args if isinstance(arg, torch.Tensor)), None) + if tensor is None: + tensor = next( + (value for value in kwargs.values() if isinstance(value, torch.Tensor)), + None, + ) + if tensor is None: + return fn(*args, **kwargs) + with torch_device_fn.device(tensor.device): + return fn(*args, **kwargs) + + +def indexer_k_quant_and_cache_flaggems(*args, **kwargs): + return _native("indexer_k_quant_and_cache", *args, **kwargs) + + +def cp_gather_indexer_k_quant_cache_flaggems(*args, **kwargs): + return _native("cp_gather_indexer_k_quant_cache", *args, **kwargs) + + +def top_k_per_row_prefill_flaggems(*args, **kwargs): + return _native("top_k_per_row_prefill", *args, **kwargs) + + +def top_k_per_row_decode_flaggems(*args, **kwargs): + return _native("top_k_per_row_decode", *args, **kwargs) + + +def pack_seq_triton_flaggems(*args, **kwargs): + return _native("pack_seq_triton", *args, **kwargs) + + +def unpack_seq_triton_flaggems(*args, **kwargs): + return _native("unpack_seq_triton", *args, **kwargs) diff --git a/vllm_fl/dispatch/backends/flaggems/register_ops.py b/vllm_fl/dispatch/backends/flaggems/register_ops.py index 24da11fbd..8e52c5cd0 100644 --- a/vllm_fl/dispatch/backends/flaggems/register_ops.py +++ b/vllm_fl/dispatch/backends/flaggems/register_ops.py @@ -39,6 +39,71 @@ def register_builtins(registry) -> None: is_avail = backend.is_available impls = [ + # DeepSeek-V4 + OpImpl( + op_name="deepseek_v4_inv_rope_quant_int8", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available( + backend.deepseek_v4_inv_rope_quant_int8, + is_avail, + ), + vendor=None, + priority=BackendPriority.DEFAULT, + ), + *[ + OpImpl( + op_name=f"deepseek_v4_{op_name}", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available( + getattr(backend, f"deepseek_v4_{op_name}"), + is_avail, + ), + vendor=None, + priority=BackendPriority.DEFAULT, + ) + for op_name in ( + "inv_rope_quant_fp8", + "int8_scaled_mm", + "mhc_pre", + "mhc_fused_post_pre", + "mhc_post", + "hc_head", + "fused_q_kv_rmsnorm", + "qnorm_rope_kv_quant_insert", + "qnorm_rope_kv_bf16_insert", + "qnorm_rope_kv_fp8_insert", + "compute_global_topk_indices_and_lens", + "flash_mla_with_kvcache", + "dequantize_and_gather_k_cache", + "combine_topk_swa_indices", + "flash_mla_sparse_fwd", + "fused_indexer_q_rope_quant", + "fused_indexer_q_rope_quant_int8", + "compress_int8_indexer_k_cache", + "int8_mqa_logits", + "int8_paged_mqa_logits", + ) + ], + *[ + OpImpl( + op_name=op_name, + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available(getattr(backend, op_name), is_avail), + vendor=None, + priority=BackendPriority.DEFAULT, + ) + for op_name in ( + "indexer_k_quant_and_cache", + "cp_gather_indexer_k_quant_cache", + "top_k_per_row_prefill", + "top_k_per_row_decode", + "pack_seq_triton", + "unpack_seq_triton", + ) + ], # Quantization OpImpl( op_name="dynamic_per_token_quant_int8", diff --git a/vllm_fl/dispatch/backends/reference/impl/deepseek_v4.py b/vllm_fl/dispatch/backends/reference/impl/deepseek_v4.py new file mode 100644 index 000000000..c703bf3da --- /dev/null +++ b/vllm_fl/dispatch/backends/reference/impl/deepseek_v4.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""PyTorch reference implementations of DeepSeek-V4-specific operators.""" + +from __future__ import annotations + +import torch + + +def deepseek_v4_fused_q_kv_rmsnorm_torch(*args, **kwargs): + from vllm.models.deepseek_v4.common.ops import fused_q_kv_rmsnorm + + return fused_q_kv_rmsnorm(*args, **kwargs) + + +def deepseek_v4_qnorm_rope_kv_quant_insert_torch(*args, **kwargs): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + *args, **kwargs + ) + + +def deepseek_v4_qnorm_rope_kv_bf16_insert_torch(*args, **kwargs): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + *args, **kwargs + ) + + +def deepseek_v4_qnorm_rope_kv_fp8_insert_torch(*args, **kwargs): + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + *args, **kwargs + ) + + +def deepseek_v4_compute_global_topk_indices_and_lens_torch(*args, **kwargs): + from vllm.models.deepseek_v4.common.ops import ( + compute_global_topk_indices_and_lens, + ) + + return compute_global_topk_indices_and_lens(*args, **kwargs) + + +def deepseek_v4_flash_mla_with_kvcache_torch(*args, **kwargs): + from vllm.v1.attention.ops.flashmla import flash_mla_with_kvcache + + return flash_mla_with_kvcache(*args, **kwargs) + + +def deepseek_v4_dequantize_and_gather_k_cache_torch(*args, **kwargs): + from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache + + return dequantize_and_gather_k_cache(*args, **kwargs) + + +def deepseek_v4_combine_topk_swa_indices_torch(*args, **kwargs): + from vllm.models.deepseek_v4.common.ops import combine_topk_swa_indices + + return combine_topk_swa_indices(*args, **kwargs) + + +def deepseek_v4_flash_mla_sparse_fwd_torch(*args, **kwargs): + from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd + + return flash_mla_sparse_fwd(*args, **kwargs) + + +def deepseek_v4_fused_indexer_q_rope_quant_torch(*args, **kwargs): + from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant + + return fused_indexer_q_rope_quant(*args, **kwargs) + + +def _int8_indexer_op(name, *args, **kwargs): + from vllm_fl.ops import deepseek_v4_int8_indexer + + return getattr(deepseek_v4_int8_indexer, name)(*args, **kwargs) + + +def deepseek_v4_fused_indexer_q_rope_quant_int8_torch(*args, **kwargs): + return _int8_indexer_op("fused_indexer_q_rope_quant_int8", *args, **kwargs) + + +def deepseek_v4_compress_int8_indexer_k_cache_torch(*args, **kwargs): + return _int8_indexer_op("compress_int8_indexer_k_cache", *args, **kwargs) + + +def deepseek_v4_int8_mqa_logits_torch(*args, **kwargs): + return _int8_indexer_op("int8_mqa_logits", *args, **kwargs) + + +def deepseek_v4_int8_paged_mqa_logits_torch(*args, **kwargs): + return _int8_indexer_op("int8_paged_mqa_logits", *args, **kwargs) + + +def deepseek_v4_inv_rope_quant_int8_torch( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply inverse RoPE and symmetric per-group-row INT8 quantization.""" + if o.ndim != 3: + raise ValueError("o must be a 3D [tokens, heads, head_dim] tensor") + tokens, heads, head_dim = o.shape + if heads != n_groups * heads_per_group: + raise ValueError("heads must equal n_groups * heads_per_group") + if head_dim != nope_dim + rope_dim: + raise ValueError("head_dim must equal nope_dim + rope_dim") + if rope_dim % 2 != 0: + raise ValueError("rope_dim must be even") + + selected_cache = cos_sin_cache.index_select(0, positions.to(torch.long)) + half_rope = rope_dim // 2 + cos = selected_cache[:, :half_rope].unsqueeze(1).to(torch.float32) + sin = selected_cache[:, half_rope:rope_dim].unsqueeze(1).to(torch.float32) + + values = o.to(torch.float32) + nope = values[..., :nope_dim] + rope = values[..., nope_dim:] + even = rope[..., 0::2] + odd = rope[..., 1::2] + inv_rope = torch.stack( + (even * cos + odd * sin, odd * cos - even * sin), + dim=-1, + ).flatten(-2) + restored = torch.cat((nope, inv_rope), dim=-1).to(torch.bfloat16) + + group_dim = heads_per_group * head_dim + grouped = ( + restored.reshape(tokens, n_groups, group_dim) + .permute(1, 0, 2) + .contiguous() + .to(torch.float32) + ) + absmax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(1.0e-4) + scales = absmax / 127.0 + normalized = grouped / scales + rounded = torch.where( + normalized >= 0, + torch.floor(normalized + 0.5), + torch.ceil(normalized - 0.5), + ) + quantized = rounded.clamp(-127, 127).to(torch.int8) + return quantized, scales + + +def deepseek_v4_int8_scaled_mm_torch(x, weight, scale_a, scale_b, out_dtype, bias=None): + out = (x.float() @ weight.float()) * scale_a.float() + out = out * scale_b.float() + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +def _rms_norm(x, weight, eps): + out = x.float() * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps) + return (out * weight.float()).to(x.dtype) + + +def deepseek_v4_mhc_pre_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits=1, + norm_weight=None, + norm_eps=1e-6, +): + from vllm.model_executor.kernels.mhc.torch import mhc_pre_torch + + post, comb, layer_input = mhc_pre_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + ) + if norm_weight is not None: + layer_input = _rms_norm(layer_input, norm_weight, norm_eps) + return post, comb, layer_input + + +def deepseek_v4_mhc_post_torch(x, residual, post_mix, comb_mix): + from vllm.model_executor.kernels.mhc.torch import mhc_post_torch + + return mhc_post_torch(x, residual, post_mix, comb_mix) + + +def deepseek_v4_mhc_fused_post_pre_torch( + x, + residual, + post_mix, + comb_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits=1, + tile_n=1, + norm_weight=None, + norm_eps=1e-6, +): + del tile_n + residual_cur = deepseek_v4_mhc_post_torch(x, residual, post_mix, comb_mix) + post_cur, comb_cur, layer_input = deepseek_v4_mhc_pre_torch( + residual_cur, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) + return residual_cur, post_cur, comb_cur, layer_input + + +def deepseek_v4_hc_head_torch(hs_flat, fn, hc_scale, hc_base, rms_eps, hc_eps): + x = hs_flat.flatten(-2).float() + x = x * torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + rms_eps) + pre = torch.sigmoid(torch.nn.functional.linear(x, fn) * hc_scale + hc_base) + pre = pre + hc_eps + return torch.sum(pre.unsqueeze(-1) * hs_flat.float(), dim=-2).to(torch.bfloat16) + + +def deepseek_v4_inv_rope_quant_fp8_torch( + o, + positions, + cos_sin_cache, + heads_per_group, + quant_group_size, + chunks_per_head, + rope_start, + half_rope, + tma_aligned_scales, + fp8_max, + tma_aligned_T, + num_tokens, + n_groups, + d, + scale_inner, +): + del chunks_per_head + if tma_aligned_scales: + raise NotImplementedError( + "reference FP8 inverse-RoPE does not pack UE8M0 scales" + ) + cache = cos_sin_cache.index_select(0, positions.long()) + values = o.float() + rope = values[..., rope_start : rope_start + 2 * half_rope] + even, odd = rope[..., 0::2], rope[..., 1::2] + cos = cache[:, :half_rope].unsqueeze(1).float() + sin = cache[:, half_rope : 2 * half_rope].unsqueeze(1).float() + restored_rope = torch.stack( + (even * cos + odd * sin, odd * cos - even * sin), dim=-1 + ).flatten(-2) + restored = torch.cat( + ( + values[..., :rope_start], + restored_rope, + values[..., rope_start + 2 * half_rope :], + ), + dim=-1, + ).to(torch.bfloat16) + grouped = restored.reshape(num_tokens, n_groups, d).permute(1, 0, 2) + padded = torch.nn.functional.pad( + grouped.float(), + (0, scale_inner * quant_group_size - d), + ).reshape(n_groups, num_tokens, scale_inner, quant_group_size) + scales = padded.abs().amax(dim=-1).clamp_min(1e-10) / fp8_max + expanded = scales.repeat_interleave(quant_group_size, dim=-1)[..., :d] + fp8_buf = ( + (grouped.float() / expanded).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + ) + scale_buf = torch.empty( + n_groups * scale_inner * tma_aligned_T, + dtype=torch.float32, + device=o.device, + ).as_strided( + (n_groups, num_tokens, scale_inner), + (scale_inner * tma_aligned_T, 1, tma_aligned_T), + ) + scale_buf.copy_(scales) + return fp8_buf, scale_buf diff --git a/vllm_fl/dispatch/backends/reference/impl/sparse_attn_indexer.py b/vllm_fl/dispatch/backends/reference/impl/sparse_attn_indexer.py new file mode 100644 index 000000000..b71c16c4f --- /dev/null +++ b/vllm_fl/dispatch/backends/reference/impl/sparse_attn_indexer.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""vLLM-native fallback implementations for sparse indexer helper ops.""" + +from __future__ import annotations + + +def indexer_k_quant_and_cache_torch(*args, **kwargs): + from vllm import _custom_ops as ops + + return ops.indexer_k_quant_and_cache(*args, **kwargs) + + +def cp_gather_indexer_k_quant_cache_torch(*args, **kwargs): + from vllm import _custom_ops as ops + + return ops.cp_gather_indexer_k_quant_cache(*args, **kwargs) + + +def top_k_per_row_prefill_torch(*args, **kwargs): + from vllm import _custom_ops as ops + + return ops.top_k_per_row_prefill(*args, **kwargs) + + +def top_k_per_row_decode_torch(*args, **kwargs): + from vllm import _custom_ops as ops + + return ops.top_k_per_row_decode(*args, **kwargs) + + +def pack_seq_triton_torch(*args, **kwargs): + from vllm.v1.attention.ops.common import pack_seq_triton + + return pack_seq_triton(*args, **kwargs) + + +def unpack_seq_triton_torch(*args, **kwargs): + from vllm.v1.attention.ops.common import unpack_seq_triton + + return unpack_seq_triton(*args, **kwargs) diff --git a/vllm_fl/dispatch/backends/reference/reference.py b/vllm_fl/dispatch/backends/reference/reference.py index 015dd2afc..5eee047c2 100644 --- a/vllm_fl/dispatch/backends/reference/reference.py +++ b/vllm_fl/dispatch/backends/reference/reference.py @@ -10,8 +10,6 @@ from __future__ import annotations -from typing import Optional, Union - import torch from vllm_fl.dispatch.backends.base import Backend @@ -25,7 +23,7 @@ class ReferenceBackend(Backend): implementations that are always available as fallbacks. """ - _available: Optional[bool] = None + _available: bool | None = None @property def name(self) -> str: @@ -35,7 +33,7 @@ def is_available(self) -> bool: """Check if PyTorch is available.""" if ReferenceBackend._available is None: try: - import torch + import torch # noqa: F401 ReferenceBackend._available = True except ImportError: @@ -43,6 +41,125 @@ def is_available(self) -> bool: return ReferenceBackend._available # ==================== Operator Implementations ==================== + def deepseek_v4_inv_rope_quant_int8( + self, + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .impl.deepseek_v4 import deepseek_v4_inv_rope_quant_int8_torch + + return deepseek_v4_inv_rope_quant_int8_torch( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + + def _deepseek_v4_call(self, op_name, *args, **kwargs): + from .impl import deepseek_v4 + + fn = getattr(deepseek_v4, f"deepseek_v4_{op_name}_torch") + return fn(*args, **kwargs) + + def deepseek_v4_inv_rope_quant_fp8(self, *args, **kwargs): + return self._deepseek_v4_call("inv_rope_quant_fp8", *args, **kwargs) + + def deepseek_v4_int8_scaled_mm(self, *args, **kwargs): + return self._deepseek_v4_call("int8_scaled_mm", *args, **kwargs) + + def deepseek_v4_mhc_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_pre", *args, **kwargs) + + def deepseek_v4_mhc_fused_post_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_fused_post_pre", *args, **kwargs) + + def deepseek_v4_mhc_post(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_post", *args, **kwargs) + + def deepseek_v4_hc_head(self, *args, **kwargs): + return self._deepseek_v4_call("hc_head", *args, **kwargs) + + def deepseek_v4_fused_q_kv_rmsnorm(self, *args, **kwargs): + return self._deepseek_v4_call("fused_q_kv_rmsnorm", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_quant_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_quant_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_bf16_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_bf16_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_fp8_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_fp8_insert", *args, **kwargs) + + def deepseek_v4_compute_global_topk_indices_and_lens(self, *args, **kwargs): + return self._deepseek_v4_call( + "compute_global_topk_indices_and_lens", *args, **kwargs + ) + + def deepseek_v4_flash_mla_with_kvcache(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_with_kvcache", *args, **kwargs) + + def deepseek_v4_dequantize_and_gather_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call("dequantize_and_gather_k_cache", *args, **kwargs) + + def deepseek_v4_combine_topk_swa_indices(self, *args, **kwargs): + return self._deepseek_v4_call("combine_topk_swa_indices", *args, **kwargs) + + def deepseek_v4_flash_mla_sparse_fwd(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_sparse_fwd", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant(self, *args, **kwargs): + return self._deepseek_v4_call("fused_indexer_q_rope_quant", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant_int8(self, *args, **kwargs): + return self._deepseek_v4_call( + "fused_indexer_q_rope_quant_int8", *args, **kwargs + ) + + def deepseek_v4_compress_int8_indexer_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call( + "compress_int8_indexer_k_cache", *args, **kwargs + ) + + def deepseek_v4_int8_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_mqa_logits", *args, **kwargs) + + def deepseek_v4_int8_paged_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_paged_mqa_logits", *args, **kwargs) + + def _sparse_indexer_call(self, op_name, *args, **kwargs): + from .impl import sparse_attn_indexer + + fn = getattr(sparse_attn_indexer, f"{op_name}_torch") + return fn(*args, **kwargs) + + def indexer_k_quant_and_cache(self, *args, **kwargs): + return self._sparse_indexer_call("indexer_k_quant_and_cache", *args, **kwargs) + + def cp_gather_indexer_k_quant_cache(self, *args, **kwargs): + return self._sparse_indexer_call( + "cp_gather_indexer_k_quant_cache", *args, **kwargs + ) + + def top_k_per_row_prefill(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_prefill", *args, **kwargs) + + def top_k_per_row_decode(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_decode", *args, **kwargs) + + def pack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("pack_seq_triton", *args, **kwargs) + + def unpack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("unpack_seq_triton", *args, **kwargs) def dynamic_per_token_quant_int8( self, @@ -86,8 +203,8 @@ def rms_norm( self, obj, x: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ RMS normalization. diff --git a/vllm_fl/dispatch/backends/reference/register_ops.py b/vllm_fl/dispatch/backends/reference/register_ops.py index 8eaf19f83..616f4081c 100644 --- a/vllm_fl/dispatch/backends/reference/register_ops.py +++ b/vllm_fl/dispatch/backends/reference/register_ops.py @@ -37,6 +37,71 @@ def register_builtins(registry) -> None: is_avail = backend.is_available impls = [ + # DeepSeek-V4 + OpImpl( + op_name="deepseek_v4_inv_rope_quant_int8", + impl_id="reference.torch", + kind=BackendImplKind.REFERENCE, + fn=_bind_is_available( + backend.deepseek_v4_inv_rope_quant_int8, + is_avail, + ), + vendor=None, + priority=BackendPriority.REFERENCE, + ), + *[ + OpImpl( + op_name=f"deepseek_v4_{op_name}", + impl_id="reference.torch", + kind=BackendImplKind.REFERENCE, + fn=_bind_is_available( + getattr(backend, f"deepseek_v4_{op_name}"), + is_avail, + ), + vendor=None, + priority=BackendPriority.REFERENCE, + ) + for op_name in ( + "inv_rope_quant_fp8", + "int8_scaled_mm", + "mhc_pre", + "mhc_fused_post_pre", + "mhc_post", + "hc_head", + "fused_q_kv_rmsnorm", + "qnorm_rope_kv_quant_insert", + "qnorm_rope_kv_bf16_insert", + "qnorm_rope_kv_fp8_insert", + "compute_global_topk_indices_and_lens", + "flash_mla_with_kvcache", + "dequantize_and_gather_k_cache", + "combine_topk_swa_indices", + "flash_mla_sparse_fwd", + "fused_indexer_q_rope_quant", + "fused_indexer_q_rope_quant_int8", + "compress_int8_indexer_k_cache", + "int8_mqa_logits", + "int8_paged_mqa_logits", + ) + ], + *[ + OpImpl( + op_name=op_name, + impl_id="reference.torch", + kind=BackendImplKind.REFERENCE, + fn=_bind_is_available(getattr(backend, op_name), is_avail), + vendor=None, + priority=BackendPriority.REFERENCE, + ) + for op_name in ( + "indexer_k_quant_and_cache", + "cp_gather_indexer_k_quant_cache", + "top_k_per_row_prefill", + "top_k_per_row_decode", + "pack_seq_triton", + "unpack_seq_triton", + ) + ], # Quantization OpImpl( op_name="dynamic_per_token_quant_int8", @@ -93,33 +158,6 @@ def register_builtins(registry) -> None: vendor=None, priority=BackendPriority.REFERENCE, ), - # MoE align - OpImpl( - op_name="moe_align_block_size", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.moe_align_block_size, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # MoE sum - OpImpl( - op_name="moe_sum", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.moe_sum, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), - # topk softmax - OpImpl( - op_name="topk_softmax", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.topk_softmax, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), # invoke fused moe triton kernel OpImpl( op_name="invoke_fused_moe_triton_kernel", @@ -129,15 +167,6 @@ def register_builtins(registry) -> None: vendor=None, priority=BackendPriority.REFERENCE, ), - # grouped topk - OpImpl( - op_name="grouped_topk", - impl_id="reference.torch", - kind=BackendImplKind.REFERENCE, - fn=_bind_is_available(backend.grouped_topk, is_avail), - vendor=None, - priority=BackendPriority.REFERENCE, - ), ] registry.register_many(impls) diff --git a/vllm_fl/dispatch/backends/vendor/cuda/cuda.py b/vllm_fl/dispatch/backends/vendor/cuda/cuda.py index b22aa5007..397586ddb 100644 --- a/vllm_fl/dispatch/backends/vendor/cuda/cuda.py +++ b/vllm_fl/dispatch/backends/vendor/cuda/cuda.py @@ -8,8 +8,6 @@ from __future__ import annotations -from typing import Optional, Union - import torch from vllm_fl.dispatch.backends.base import Backend @@ -23,23 +21,24 @@ class CudaBackend(Backend): operator implementations for NVIDIA GPUs. """ - _available: Optional[bool] = None + _available: bool | None = None @property def name(self) -> str: return "cuda" @property - def vendor(self) -> Optional[str]: + def vendor(self) -> str | None: return "nvidia" def is_available(self) -> bool: """ Check if CUDA hardware and libraries are available. - This method uses the platform's vendor information from FlagGems - to determine if the device is a real NVIDIA GPU, decoupling from - CUDA-alike devices (MACA, MUSA, etc.) which have their own vendor names. + Use vLLM's platform discriminator rather than ``device_name``. The + in-tree NVIDIA platform reports ``device_name == "cuda"`` and has no + ``vendor_name``, while CUDA-alike out-of-tree platforms override + ``is_cuda()`` to remain distinct from NVIDIA. """ if CudaBackend._available is None: try: @@ -50,18 +49,131 @@ def is_available(self) -> bool: from vllm.platforms import current_platform - if ( - hasattr(current_platform, "device_name") - and current_platform.device_name == "nvidia" - ): - CudaBackend._available = True - else: - CudaBackend._available = False + CudaBackend._available = current_platform.is_cuda() except Exception: CudaBackend._available = False return CudaBackend._available # ==================== Operator Implementations ==================== + def deepseek_v4_inv_rope_quant_int8( + self, + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .impl.deepseek_v4 import deepseek_v4_inv_rope_quant_int8_cuda + + return deepseek_v4_inv_rope_quant_int8_cuda( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + + def _deepseek_v4_call(self, op_name, *args, **kwargs): + from .impl import deepseek_v4 + + fn = getattr(deepseek_v4, f"deepseek_v4_{op_name}_cuda") + return fn(*args, **kwargs) + + def deepseek_v4_inv_rope_quant_fp8(self, *args, **kwargs): + return self._deepseek_v4_call("inv_rope_quant_fp8", *args, **kwargs) + + def deepseek_v4_int8_scaled_mm(self, *args, **kwargs): + return self._deepseek_v4_call("int8_scaled_mm", *args, **kwargs) + + def deepseek_v4_mhc_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_pre", *args, **kwargs) + + def deepseek_v4_mhc_fused_post_pre(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_fused_post_pre", *args, **kwargs) + + def deepseek_v4_mhc_post(self, *args, **kwargs): + return self._deepseek_v4_call("mhc_post", *args, **kwargs) + + def deepseek_v4_hc_head(self, *args, **kwargs): + return self._deepseek_v4_call("hc_head", *args, **kwargs) + + def deepseek_v4_fused_q_kv_rmsnorm(self, *args, **kwargs): + return self._deepseek_v4_call("fused_q_kv_rmsnorm", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_quant_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_quant_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_bf16_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_bf16_insert", *args, **kwargs) + + def deepseek_v4_qnorm_rope_kv_fp8_insert(self, *args, **kwargs): + return self._deepseek_v4_call("qnorm_rope_kv_fp8_insert", *args, **kwargs) + + def deepseek_v4_compute_global_topk_indices_and_lens(self, *args, **kwargs): + return self._deepseek_v4_call( + "compute_global_topk_indices_and_lens", *args, **kwargs + ) + + def deepseek_v4_flash_mla_with_kvcache(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_with_kvcache", *args, **kwargs) + + def deepseek_v4_dequantize_and_gather_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call("dequantize_and_gather_k_cache", *args, **kwargs) + + def deepseek_v4_combine_topk_swa_indices(self, *args, **kwargs): + return self._deepseek_v4_call("combine_topk_swa_indices", *args, **kwargs) + + def deepseek_v4_flash_mla_sparse_fwd(self, *args, **kwargs): + return self._deepseek_v4_call("flash_mla_sparse_fwd", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant(self, *args, **kwargs): + return self._deepseek_v4_call("fused_indexer_q_rope_quant", *args, **kwargs) + + def deepseek_v4_fused_indexer_q_rope_quant_int8(self, *args, **kwargs): + return self._deepseek_v4_call( + "fused_indexer_q_rope_quant_int8", *args, **kwargs + ) + + def deepseek_v4_compress_int8_indexer_k_cache(self, *args, **kwargs): + return self._deepseek_v4_call( + "compress_int8_indexer_k_cache", *args, **kwargs + ) + + def deepseek_v4_int8_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_mqa_logits", *args, **kwargs) + + def deepseek_v4_int8_paged_mqa_logits(self, *args, **kwargs): + return self._deepseek_v4_call("int8_paged_mqa_logits", *args, **kwargs) + + def _sparse_indexer_call(self, op_name, *args, **kwargs): + from .impl import sparse_attn_indexer + + fn = getattr(sparse_attn_indexer, f"{op_name}_cuda") + return fn(*args, **kwargs) + + def indexer_k_quant_and_cache(self, *args, **kwargs): + return self._sparse_indexer_call("indexer_k_quant_and_cache", *args, **kwargs) + + def cp_gather_indexer_k_quant_cache(self, *args, **kwargs): + return self._sparse_indexer_call( + "cp_gather_indexer_k_quant_cache", *args, **kwargs + ) + + def top_k_per_row_prefill(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_prefill", *args, **kwargs) + + def top_k_per_row_decode(self, *args, **kwargs): + return self._sparse_indexer_call("top_k_per_row_decode", *args, **kwargs) + + def pack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("pack_seq_triton", *args, **kwargs) + + def unpack_seq_triton(self, *args, **kwargs): + return self._sparse_indexer_call("unpack_seq_triton", *args, **kwargs) def silu_and_mul(self, obj, x: torch.Tensor) -> torch.Tensor: """ @@ -101,8 +213,8 @@ def rms_norm( self, obj, x: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ RMS normalization using vLLM's CUDA implementation. @@ -189,7 +301,7 @@ def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int, - expert_map: Optional[torch.Tensor] = None, + expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ignore_invalid_experts: bool = False, ): @@ -285,6 +397,12 @@ def grouped_topk( from .impl.fused_moe import grouped_topk_cuda return grouped_topk_cuda( - scores, n_group, topk_group, topk, - renormalize, routed_scaling_factor, bias, scoring_func, + scores, + n_group, + topk_group, + topk, + renormalize, + routed_scaling_factor, + bias, + scoring_func, ) diff --git a/vllm_fl/dispatch/backends/vendor/cuda/impl/deepseek_v4.py b/vllm_fl/dispatch/backends/vendor/cuda/impl/deepseek_v4.py new file mode 100644 index 000000000..ad3491ebc --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/cuda/impl/deepseek_v4.py @@ -0,0 +1,189 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""CUDA vendor implementation of DeepSeek-V4-specific operators.""" + +from __future__ import annotations + +import torch + +from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( + _fused_inv_rope_fp8_quant_per_head, +) + +from vllm_fl.ops.deepseek_v4_int8_woa import ( + fused_inv_rope_quant_int8_triton, +) + + +def _vllm_dsv4_op(op_name, *args, **kwargs): + """Call the vLLM CUDA implementation behind an OpManager vendor entry.""" + from vllm_fl.dispatch.backends.reference.impl import deepseek_v4 + + fn = getattr(deepseek_v4, f"deepseek_v4_{op_name}_torch") + return fn(*args, **kwargs) + + +def deepseek_v4_fused_q_kv_rmsnorm_cuda(*args, **kwargs): + return _vllm_dsv4_op("fused_q_kv_rmsnorm", *args, **kwargs) + + +def deepseek_v4_qnorm_rope_kv_quant_insert_cuda(*args, **kwargs): + return _vllm_dsv4_op("qnorm_rope_kv_quant_insert", *args, **kwargs) + + +def deepseek_v4_qnorm_rope_kv_bf16_insert_cuda(*args, **kwargs): + return _vllm_dsv4_op("qnorm_rope_kv_bf16_insert", *args, **kwargs) + + +def deepseek_v4_qnorm_rope_kv_fp8_insert_cuda(*args, **kwargs): + return _vllm_dsv4_op("qnorm_rope_kv_fp8_insert", *args, **kwargs) + + +def deepseek_v4_compute_global_topk_indices_and_lens_cuda(*args, **kwargs): + return _vllm_dsv4_op("compute_global_topk_indices_and_lens", *args, **kwargs) + + +def deepseek_v4_flash_mla_with_kvcache_cuda(*args, **kwargs): + return _vllm_dsv4_op("flash_mla_with_kvcache", *args, **kwargs) + + +def deepseek_v4_dequantize_and_gather_k_cache_cuda(*args, **kwargs): + return _vllm_dsv4_op("dequantize_and_gather_k_cache", *args, **kwargs) + + +def deepseek_v4_combine_topk_swa_indices_cuda(*args, **kwargs): + return _vllm_dsv4_op("combine_topk_swa_indices", *args, **kwargs) + + +def deepseek_v4_flash_mla_sparse_fwd_cuda(*args, **kwargs): + return _vllm_dsv4_op("flash_mla_sparse_fwd", *args, **kwargs) + + +def deepseek_v4_fused_indexer_q_rope_quant_cuda(*args, **kwargs): + return _vllm_dsv4_op("fused_indexer_q_rope_quant", *args, **kwargs) + + +def deepseek_v4_fused_indexer_q_rope_quant_int8_cuda(*args, **kwargs): + return _vllm_dsv4_op("fused_indexer_q_rope_quant_int8", *args, **kwargs) + + +def deepseek_v4_compress_int8_indexer_k_cache_cuda(*args, **kwargs): + return _vllm_dsv4_op("compress_int8_indexer_k_cache", *args, **kwargs) + + +def deepseek_v4_int8_mqa_logits_cuda(*args, **kwargs): + return _vllm_dsv4_op("int8_mqa_logits", *args, **kwargs) + + +def deepseek_v4_int8_paged_mqa_logits_cuda(*args, **kwargs): + return _vllm_dsv4_op("int8_paged_mqa_logits", *args, **kwargs) + + +def deepseek_v4_inv_rope_quant_int8_cuda( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the fused Triton implementation on NVIDIA CUDA devices.""" + return fused_inv_rope_quant_int8_triton( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + + +def deepseek_v4_inv_rope_quant_fp8_cuda( + o, + positions, + cos_sin_cache, + heads_per_group, + quant_group_size, + chunks_per_head, + rope_start, + half_rope, + tma_aligned_scales, + fp8_max, + tma_aligned_T, + num_tokens, + n_groups, + d, + scale_inner, +): + fp8_buf = torch.empty( + (n_groups, num_tokens, d), dtype=torch.float8_e4m3fn, device=o.device + ) + scale_dtype = torch.int32 if tma_aligned_scales else torch.float32 + scale_buf = torch.empty( + n_groups * scale_inner * tma_aligned_T, + dtype=scale_dtype, + device=o.device, + ).as_strided( + (n_groups, num_tokens, scale_inner), + (scale_inner * tma_aligned_T, 1, tma_aligned_T), + ) + grid = (tma_aligned_T, n_groups * heads_per_group) + _fused_inv_rope_fp8_quant_per_head[grid]( + o, + positions, + cos_sin_cache, + fp8_buf, + scale_buf, + num_tokens, + heads_per_group=heads_per_group, + o_stride_token=o.stride(0), + o_stride_head=o.stride(1), + cache_stride_pos=cos_sin_cache.stride(0), + fp8_stride_group=fp8_buf.stride(0), + fp8_stride_token=fp8_buf.stride(1), + scale_stride_group=scale_buf.stride(0), + scale_stride_k=scale_buf.stride(2), + fp8_max=fp8_max, + eps=1e-10, + QUANT_GROUP_SIZE=quant_group_size, + CHUNKS_PER_HEAD=chunks_per_head, + ROPE_START=rope_start, + HALF_ROPE=half_rope, + TMA_ALIGNED_SCALES=tma_aligned_scales, + USE_GDC=False, + launch_pdl=False, + num_stages=1, + num_warps=1, + ) + return fp8_buf, scale_buf + + +def deepseek_v4_int8_scaled_mm_cuda(x, weight, scale_a, scale_b, out_dtype, bias=None): + from vllm import _custom_ops as ops + + return ops.cutlass_scaled_mm( + x, + weight, + scale_a=scale_a, + scale_b=scale_b, + bias=bias, + out_dtype=out_dtype, + ) + + +def deepseek_v4_mhc_pre_cuda(*args): + return torch.ops.vllm.mhc_pre_tilelang(*args) + + +def deepseek_v4_mhc_fused_post_pre_cuda(*args): + return torch.ops.vllm.mhc_fused_post_pre_tilelang(*args) + + +def deepseek_v4_mhc_post_cuda(*args): + return torch.ops.vllm.mhc_post_tilelang(*args) + + +def deepseek_v4_hc_head_cuda(*args): + return torch.ops.vllm.hc_head_fused_kernel_tilelang(*args) diff --git a/vllm_fl/dispatch/backends/vendor/cuda/impl/sparse_attn_indexer.py b/vllm_fl/dispatch/backends/vendor/cuda/impl/sparse_attn_indexer.py new file mode 100644 index 000000000..7a8610054 --- /dev/null +++ b/vllm_fl/dispatch/backends/vendor/cuda/impl/sparse_attn_indexer.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 BAAI. All rights reserved. + +"""CUDA vendor implementations for sparse indexer helper ops.""" + +from __future__ import annotations + + +def _native(op_name, *args, **kwargs): + from vllm_fl.dispatch.backends.reference.impl import sparse_attn_indexer + + fn = getattr(sparse_attn_indexer, f"{op_name}_torch") + return fn(*args, **kwargs) + + +def indexer_k_quant_and_cache_cuda(*args, **kwargs): + return _native("indexer_k_quant_and_cache", *args, **kwargs) + + +def cp_gather_indexer_k_quant_cache_cuda(*args, **kwargs): + return _native("cp_gather_indexer_k_quant_cache", *args, **kwargs) + + +def top_k_per_row_prefill_cuda(*args, **kwargs): + return _native("top_k_per_row_prefill", *args, **kwargs) + + +def top_k_per_row_decode_cuda(*args, **kwargs): + return _native("top_k_per_row_decode", *args, **kwargs) + + +def pack_seq_triton_cuda(*args, **kwargs): + return _native("pack_seq_triton", *args, **kwargs) + + +def unpack_seq_triton_cuda(*args, **kwargs): + return _native("unpack_seq_triton", *args, **kwargs) diff --git a/vllm_fl/dispatch/backends/vendor/cuda/register_ops.py b/vllm_fl/dispatch/backends/vendor/cuda/register_ops.py index 6a9a16425..228a86d08 100644 --- a/vllm_fl/dispatch/backends/vendor/cuda/register_ops.py +++ b/vllm_fl/dispatch/backends/vendor/cuda/register_ops.py @@ -10,7 +10,7 @@ import functools -from vllm_fl.dispatch.types import OpImpl, BackendImplKind, BackendPriority +from vllm_fl.dispatch.types import BackendImplKind, BackendPriority, OpImpl def _bind_is_available(fn, is_available_fn): @@ -37,6 +37,71 @@ def register_builtins(registry) -> None: is_avail = backend.is_available impls = [ + # DeepSeek-V4 + OpImpl( + op_name="deepseek_v4_inv_rope_quant_int8", + impl_id="vendor.cuda", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available( + backend.deepseek_v4_inv_rope_quant_int8, + is_avail, + ), + vendor="cuda", + priority=BackendPriority.VENDOR, + ), + *[ + OpImpl( + op_name=f"deepseek_v4_{op_name}", + impl_id="vendor.cuda", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available( + getattr(backend, f"deepseek_v4_{op_name}"), + is_avail, + ), + vendor="cuda", + priority=BackendPriority.VENDOR, + ) + for op_name in ( + "inv_rope_quant_fp8", + "int8_scaled_mm", + "mhc_pre", + "mhc_fused_post_pre", + "mhc_post", + "hc_head", + "fused_q_kv_rmsnorm", + "qnorm_rope_kv_quant_insert", + "qnorm_rope_kv_bf16_insert", + "qnorm_rope_kv_fp8_insert", + "compute_global_topk_indices_and_lens", + "flash_mla_with_kvcache", + "dequantize_and_gather_k_cache", + "combine_topk_swa_indices", + "flash_mla_sparse_fwd", + "fused_indexer_q_rope_quant", + "fused_indexer_q_rope_quant_int8", + "compress_int8_indexer_k_cache", + "int8_mqa_logits", + "int8_paged_mqa_logits", + ) + ], + *[ + OpImpl( + op_name=op_name, + impl_id="vendor.cuda", + kind=BackendImplKind.VENDOR, + fn=_bind_is_available(getattr(backend, op_name), is_avail), + vendor="cuda", + priority=BackendPriority.VENDOR, + ) + for op_name in ( + "indexer_k_quant_and_cache", + "cp_gather_indexer_k_quant_cache", + "top_k_per_row_prefill", + "top_k_per_row_decode", + "pack_seq_triton", + "unpack_seq_triton", + ) + ], # Activation OpImpl( op_name="silu_and_mul", diff --git a/vllm_fl/dispatch/builtin_ops.py b/vllm_fl/dispatch/builtin_ops.py index 47b81930f..6cec01521 100644 --- a/vllm_fl/dispatch/builtin_ops.py +++ b/vllm_fl/dispatch/builtin_ops.py @@ -46,12 +46,26 @@ def _find_vendor_backend_dir( ) -def _get_current_vendor_backend_dirs(available_vendor_dirs: set[str]) -> set[str]: - """Detect current platform vendor name and return its backend directory.""" +def _get_current_vendor_backend_dirs( + available_vendor_dirs: set[str], +) -> str | None: + """Detect the current platform and return its vendor backend directory. + + In-tree vLLM CUDA platforms (for example ``NvmlCudaPlatform``) do not + expose ``vendor_name``. ``Platform.is_cuda()`` is the canonical vLLM + discriminator for NVIDIA CUDA and, unlike matching ``device_name == + "cuda"``, does not opt CUDA-alike out-of-tree platforms into the NVIDIA + backend. + """ try: from vllm.platforms import current_platform vendor_name = getattr(current_platform, "vendor_name", None) + if ( + (not isinstance(vendor_name, str) or not vendor_name) + and current_platform.is_cuda() + ): + vendor_name = "nvidia" if not isinstance(vendor_name, str) or not vendor_name: return None return _find_vendor_backend_dir(vendor_name, available_vendor_dirs) @@ -164,4 +178,3 @@ def register_builtins(registry: OpRegistry) -> None: except Exception as e: logger.debug(f"Plugin discovery failed: {e}") - diff --git a/vllm_fl/dispatch/config/nvidia.yaml b/vllm_fl/dispatch/config/nvidia.yaml index 199a04954..ac23020b4 100644 --- a/vllm_fl/dispatch/config/nvidia.yaml +++ b/vllm_fl/dispatch/config/nvidia.yaml @@ -42,6 +42,30 @@ op_backends: - flagos - vendor - reference + indexer_k_quant_and_cache: + - flagos + - vendor:cuda + - reference + cp_gather_indexer_k_quant_cache: + - flagos + - vendor:cuda + - reference + top_k_per_row_prefill: + - flagos + - vendor:cuda + - reference + top_k_per_row_decode: + - flagos + - vendor:cuda + - reference + pack_seq_triton: + - flagos + - vendor:cuda + - reference + unpack_seq_triton: + - flagos + - vendor:cuda + - reference # FlagOS operator blacklist flagos_blacklist: diff --git a/vllm_fl/models/deepseek_v4.py b/vllm_fl/models/deepseek_v4.py new file mode 100644 index 000000000..39aa7fa5c --- /dev/null +++ b/vllm_fl/models/deepseek_v4.py @@ -0,0 +1,909 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ModelRegistry thin model adding W8A8 DSV4 output projection support.""" + +from __future__ import annotations + +from itertools import islice + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.utils import PPMissingLayer, make_layers +from vllm.models.deepseek_v4.nvidia.flashmla import ( + DeepseekV4FlashMLAAttention, +) +from vllm.models.deepseek_v4.nvidia.model import ( + DeepseekV4DecoderLayer, + DeepseekV4ForCausalLM, + DeepseekV4Model, + DeepseekV4MoE, + _select_dsv4_attn_cls, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.utils.multi_stream_utils import execute_in_parallel, maybe_execute_in_parallel +from vllm.utils.torch_utils import direct_register_custom_op, vllm_lib +from vllm.v1.worker.workspace import current_workspace_manager + +from vllm_fl.ops.deepseek_v4 import ( + combine_topk_swa_indices, + compress_int8_indexer_k_cache, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, + flash_mla_sparse_fwd, + flash_mla_with_kvcache, + fused_indexer_q_rope_quant, + fused_indexer_q_rope_quant_int8, + fused_q_kv_rmsnorm, + hc_head, + int8_scaled_mm, + inv_rope_quant_fp8, + mhc_fused_post_pre, + mhc_post, + mhc_pre, + qnorm_rope_kv_bf16_insert, + qnorm_rope_kv_fp8_insert, + qnorm_rope_kv_quant_insert, +) +from vllm_fl.ops.deepseek_v4_int8_woa import fused_inv_rope_quant_int8 +from vllm_fl.ops.sparse_attn_indexer import SparseAttnIndexerFL + + +def _patch_hopper_fp8_inv_rope_kernel() -> None: + """Bridge the vLLM/Triton PDL argument mismatch on pre-SM100 GPUs.""" + capability = current_platform.get_device_capability() + if capability is None or capability.major >= 10: + return + vllm_lib.impl( + "fused_inv_rope_fp8_quant_kernel", + inv_rope_quant_fp8, + dispatch_key=current_platform.dispatch_key, + allow_override=True, + ) + + +_patch_hopper_fp8_inv_rope_kernel() + + +def _deepseek_v4_fl_attention( + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + layer_name: str, +) -> None: + layer = get_forward_context().no_compile_layers[layer_name] + layer.attention_impl( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + out, + ) + + +def _deepseek_v4_fl_attention_fake( + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + layer_name: str, +) -> None: + del ( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + out, + layer_name, + ) + + +direct_register_custom_op( + op_name="deepseek_v4_fl_attention", + op_func=_deepseek_v4_fl_attention, + mutates_args=["out"], + fake_impl=_deepseek_v4_fl_attention_fake, +) + + +def _deepseek_v4_fl_fp8_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + hidden_size: int, + layer_name: str, +) -> torch.Tensor: + del hidden_size + layer = get_forward_context().no_compile_layers[layer_name] + return DeepseekV4FlashMLAAttention._o_proj(layer, o, positions) + + +def _deepseek_v4_fl_fp8_o_proj_fake( + o: torch.Tensor, + positions: torch.Tensor, + hidden_size: int, + layer_name: str, +) -> torch.Tensor: + del positions, layer_name + return torch.empty( + (o.shape[0], hidden_size), + dtype=o.dtype, + device=o.device, + ) + + +direct_register_custom_op( + op_name="deepseek_v4_fl_fp8_o_proj", + op_func=_deepseek_v4_fl_fp8_o_proj, + mutates_args=[], + fake_impl=_deepseek_v4_fl_fp8_o_proj_fake, +) + + +class DeepseekV4FLFlashMLAAttention(DeepseekV4FlashMLAAttention): + """FlashMLA attention with a W8A8-only wo_a branch.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + indexer = self.indexer + if indexer is None: + return + quantization_config = getattr(indexer.config, "quantization_config", None) + use_int8_kv = ( + isinstance(quantization_config, dict) + and quantization_config.get("format") == "int-quantized" + ) + indexer.use_int8_kv = use_int8_kv + if not use_int8_kv: + return + + # The 512-wide Sparse FlashMLA cache remains FP8. Only the independent + # 128-wide lightning-indexer cache switches to INT8 plus one fp32 scale. + indexer.use_fp4_kv = False + indexer.compressor.use_fp4_cache = False + native_op = indexer.indexer_op + indexer.indexer_op = SparseAttnIndexerFL( + native_op.k_cache, + native_op.quant_block_size, + native_op.scale_fmt, + native_op.topk_tokens, + native_op.head_dim, + native_op.max_model_len, + native_op.max_total_seq_len, + native_op.topk_indices_buffer, + skip_k_cache_insert=native_op.skip_k_cache_insert, + use_fp4_cache=False, + ) + + def _indexer_forward( + self, + indexer, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + compressor = indexer.compressor + + def wq_b_and_q_quant(): + q, _ = indexer.wq_b(qr) + q = q.view(-1, indexer.n_head, indexer.head_dim) + if getattr(indexer, "use_int8_kv", False): + return fused_indexer_q_rope_quant_int8( + positions, + q, + self.indexer_rotary_emb.cos_sin_cache, + indexer_weights, + indexer.softmax_scale, + indexer.n_head**-0.5, + ) + return fused_indexer_q_rope_quant( + positions, + q, + self.indexer_rotary_emb.cos_sin_cache, + indexer_weights, + indexer.softmax_scale, + indexer.n_head**-0.5, + use_fp4=indexer.use_fp4_kv, + ) + + (q_quant, weights), k = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: ( + compress_int8_indexer_k_cache( + compressor, + compressed_kv_score, + positions, + self.indexer_rotary_emb, + ) + if getattr(indexer, "use_int8_kv", False) + else compressor( + compressed_kv_score, + positions, + self.indexer_rotary_emb, + ) + ), + indexer.ln_events[0], + indexer.ln_events[1], + indexer.aux_stream, + ) + return indexer.indexer_op(hidden_states, q_quant, k, weights) + + def attention_impl( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + ) -> None: + """Upstream orchestration with every direct CUDA helper dispatched.""" + attn_metadata = get_forward_context().attn_metadata + if self.indexer is not None: + aux_streams = self.aux_stream_list + indexer = self.indexer + assert self.compressor is not None + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: self._indexer_forward( + indexer, + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + self.ln_events[0], + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, + ) + elif self.compressor is not None: + aux_stream = ( + self.aux_stream_list[0] if self.aux_stream_list is not None else None + ) + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + self.forward_mqa(q, kv, positions, out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata, + ) -> torch.Tensor: + """Route each cache-layout-specific fused insert through OpManager.""" + if not isinstance(attn_metadata, dict): + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = attn_metadata.get(self.swa_cache_layer.prefix) + assert swa_metadata is not None + swa_kv_cache = self.swa_cache_layer.kv_cache + assert positions.dtype == torch.int64 + cos_sin_cache = self.rotary_emb.cos_sin_cache + + if swa_kv_cache.dtype == torch.uint8: + return qnorm_rope_kv_quant_insert( + q, + kv, + swa_kv_cache.view(swa_kv_cache.shape[0], -1), + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + + block_size = swa_metadata.block_size + swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim) + if swa_kv_cache.dtype == torch.bfloat16: + qnorm_rope_kv_bf16_insert( + q, + kv, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.eps, + block_size, + ) + return q + + q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn) + qnorm_rope_kv_fp8_insert( + q, + kv, + q_fp8, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self._flashinfer_fp8_kv_scale, + self._flashinfer_fp8_q_scale_inv, + self.eps, + block_size, + ) + return q_fp8 + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, + swa_metadata, + attn_metadata, + swa_only: bool, + output: torch.Tensor, + ) -> None: + """Decode path with all DSV4 CUDA helpers behind OpManager.""" + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // self.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + q = q.unsqueeze(1) + swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + if self.compress_ratio <= 1: + tile_metadata = swa_metadata.tile_sched_swaonly + elif self.compress_ratio == 4: + tile_metadata = swa_metadata.tile_sched_c4a + elif self.compress_ratio == 128: + tile_metadata = swa_metadata.tile_sched_c128a + else: + raise ValueError( + f"Unsupported compress_ratio={self.compress_ratio}; " + "expected 1, 4, or 128." + ) + assert tile_metadata is not None + flash_mla_with_kvcache( + q=q, + k_cache=swa_cache, + block_table=None, + head_dim_v=512, + tile_scheduler_metadata=tile_metadata, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_metadata.decode_swa_indices, + topk_length=swa_metadata.decode_swa_lens, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + extra_k_cache=kv_cache if not swa_only else None, + extra_indices_in_kvcache=topk_indices, + extra_topk_length=topk_lens, + out=output.unsqueeze(1), + ) + + def _forward_prefill( + self, + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata, + swa_metadata, + ) -> None: + """Prefill path with gather/index construction/attention dispatched.""" + del positions + swa_only = attn_metadata is None + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert seq_lens is not None and gather_lens is not None + assert query_start_loc_cpu is not None and query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + top_k = topk_indices.shape[-1] + else: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + + chunk_plan = swa_metadata.get_prefill_chunk_plan( + compress_ratio=self.compress_ratio, + prefill_chunk_size=self.PREFILL_CHUNK_SIZE, + ) + assert chunk_plan + workspace_manager = current_workspace_manager() + for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: + chunk_size = chunk_end - chunk_start + kv = workspace_manager.get_simultaneous( + ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16), + )[0] + if not swa_only: + assert attn_metadata is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // self.compress_ratio, + offset=0, + ) + + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=chunk_n, + ) + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + chunk_m, + chunk_n, + ) + flash_mla_sparse_fwd( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices.unsqueeze(1), + sm_scale=self.scale, + attn_sink=self.attn_sink, + topk_length=combined_lens, + out=output[query_start:query_end], + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + del llama_4_scaling + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self.attn_gemm_parallel_execute(hidden_states) + ) + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + + torch.ops.vllm.deepseek_v4_fl_attention( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + self.prefix, + ) + o = o_padded[:, : self.n_local_heads, :] + return self._o_proj(o, positions) + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + grouped_weight = getattr(self.wo_a, "_fl_w8a8_grouped_weight", None) + grouped_scale = getattr(self.wo_a, "_fl_w8a8_grouped_weight_scale", None) + if grouped_weight is None or grouped_scale is None: + weight = getattr(self.wo_a, "weight", None) + weight_scale = getattr(self.wo_a, "weight_scale", None) + if weight is None or weight.dtype != torch.int8 or weight_scale is None: + return torch.ops.vllm.deepseek_v4_fl_fp8_o_proj( + o, + positions, + self.hidden_size, + self.prefix, + ) + output_per_group = weight.shape[1] // self.n_local_groups + grouped_scale = weight_scale.reshape(self.n_local_groups, output_per_group) + + o_q, o_scale = fused_inv_rope_quant_int8( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.n_local_groups, + self.n_local_heads // self.n_local_groups, + self.nope_head_dim, + self.rope_head_dim, + ) + outputs = [] + for group_idx in range(self.n_local_groups): + if grouped_weight is None: + start = group_idx * output_per_group + group_weight = weight[:, start : start + output_per_group] + else: + group_weight = grouped_weight[group_idx].transpose(0, 1) + outputs.append( + int8_scaled_mm( + o_q[group_idx], + group_weight, + o_scale[group_idx], + grouped_scale[group_idx], + o.dtype, + ) + ) + return self.wo_b(torch.stack(outputs, dim=1).flatten(1)) + + +class DeepseekV4FLDecoderLayer(DeepseekV4DecoderLayer): + """Upstream decoder layer with a thin attention-class substitution.""" + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ) -> None: + nn.Module.__init__(self) + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + attn_cls = _select_dsv4_attn_cls(vllm_config) + if attn_cls is DeepseekV4FlashMLAAttention: + attn_cls = DeepseekV4FLFlashMLAAttention + self.attn = attn_cls( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty((mix_hc, hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty((mix_hc, hc_dim), dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), + requires_grad=False, + ) + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Keep NVIDIA MHC kernels opaque to Dynamo via registered ops.""" + attn_norm_weight = self.attn_norm.weight.data + attn_norm_eps = self.attn_norm.variance_epsilon + if residual is None: + residual = x + post_mix, res_mix, x = mhc_pre( + x, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + 1, + attn_norm_weight, + attn_norm_eps, + ) + else: + assert post_mix is not None and res_mix is not None + residual, post_mix, res_mix, x = mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + 1, + 1, + attn_norm_weight, + attn_norm_eps, + ) + + x = self.attn(positions, x, None) + + ffn_norm_weight = self.ffn_norm.weight.data + ffn_norm_eps = self.ffn_norm.variance_epsilon + residual, post_mix, res_mix, x = mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + 1, + 1, + ffn_norm_weight, + ffn_norm_eps, + ) + x = self.ffn(x, input_ids) + return x, residual, post_mix, res_mix + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class DeepseekV4FLModel(DeepseekV4Model): + """Upstream DSV4 model whose layers use the FL decoder subclass.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + nn.Module.__init__(self) + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.parallel_config = vllm_config.parallel_config + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError("DeepSeek V4 MegaMoE requires expert parallel") + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + ) + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV4FLDecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + if get_pp_group().is_last_rank: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + ) + else: + self._mtp_hidden_buffer = None + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + """Run the upstream model flow with registered MHC custom ops.""" + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + if self.use_mega_moe: + input_ids = input_ids.to(torch.int64) + + residual, post_mix, res_mix = None, None, None + layer = None + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + if layer is not None: + hidden_states = mhc_post(hidden_states, residual, post_mix, res_mix) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + hidden_states = hc_head( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return self.norm(hidden_states) + + +class DeepseekV4FLForCausalLM(DeepseekV4ForCausalLM): + """Registry entry retaining upstream behavior outside INT8 wo_a.""" + + model_cls = DeepseekV4FLModel + + +__all__ = ["DeepseekV4FLForCausalLM"] diff --git a/vllm_fl/ops/deepseek_v4.py b/vllm_fl/ops/deepseek_v4.py new file mode 100644 index 000000000..c7a518480 --- /dev/null +++ b/vllm_fl/ops/deepseek_v4.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OpManager frontends for DeepSeek-V4 model-specific compute.""" + +from __future__ import annotations + +from vllm_fl.dispatch import resolve_op + +_OPS = { + # Resolve once while the model module is imported. These frontends are + # called from vLLM's full-graph compiled model; a lazy CachedOp lookup would + # make Dynamo trace OpManager's RLock on the first profile run. Import-time + # resolution keeps backend selection in OpManager while exposing only the + # selected callable to torch.compile and CUDA graph capture. + name: resolve_op(f"deepseek_v4_{name}") + for name in ( + "inv_rope_quant_fp8", + "int8_scaled_mm", + "mhc_pre", + "mhc_fused_post_pre", + "mhc_post", + "hc_head", + "fused_q_kv_rmsnorm", + "qnorm_rope_kv_quant_insert", + "qnorm_rope_kv_bf16_insert", + "qnorm_rope_kv_fp8_insert", + "compute_global_topk_indices_and_lens", + "flash_mla_with_kvcache", + "dequantize_and_gather_k_cache", + "combine_topk_swa_indices", + "flash_mla_sparse_fwd", + "fused_indexer_q_rope_quant", + "fused_indexer_q_rope_quant_int8", + "compress_int8_indexer_k_cache", + "int8_mqa_logits", + "int8_paged_mqa_logits", + ) +} + + +def inv_rope_quant_fp8(*args, **kwargs): + return _OPS["inv_rope_quant_fp8"](*args, **kwargs) + + +def int8_scaled_mm(*args, **kwargs): + return _OPS["int8_scaled_mm"](*args, **kwargs) + + +def mhc_pre(*args, **kwargs): + return _OPS["mhc_pre"](*args, **kwargs) + + +def mhc_fused_post_pre(*args, **kwargs): + return _OPS["mhc_fused_post_pre"](*args, **kwargs) + + +def mhc_post(*args, **kwargs): + return _OPS["mhc_post"](*args, **kwargs) + + +def hc_head(*args, **kwargs): + return _OPS["hc_head"](*args, **kwargs) + + +def fused_q_kv_rmsnorm(*args, **kwargs): + return _OPS["fused_q_kv_rmsnorm"](*args, **kwargs) + + +def qnorm_rope_kv_quant_insert(*args, **kwargs): + return _OPS["qnorm_rope_kv_quant_insert"](*args, **kwargs) + + +def qnorm_rope_kv_bf16_insert(*args, **kwargs): + return _OPS["qnorm_rope_kv_bf16_insert"](*args, **kwargs) + + +def qnorm_rope_kv_fp8_insert(*args, **kwargs): + return _OPS["qnorm_rope_kv_fp8_insert"](*args, **kwargs) + + +def compute_global_topk_indices_and_lens(*args, **kwargs): + return _OPS["compute_global_topk_indices_and_lens"](*args, **kwargs) + + +def flash_mla_with_kvcache(*args, **kwargs): + return _OPS["flash_mla_with_kvcache"](*args, **kwargs) + + +def dequantize_and_gather_k_cache(*args, **kwargs): + return _OPS["dequantize_and_gather_k_cache"](*args, **kwargs) + + +def combine_topk_swa_indices(*args, **kwargs): + return _OPS["combine_topk_swa_indices"](*args, **kwargs) + + +def flash_mla_sparse_fwd(*args, **kwargs): + return _OPS["flash_mla_sparse_fwd"](*args, **kwargs) + + +def fused_indexer_q_rope_quant(*args, **kwargs): + return _OPS["fused_indexer_q_rope_quant"](*args, **kwargs) + + +def fused_indexer_q_rope_quant_int8(*args, **kwargs): + return _OPS["fused_indexer_q_rope_quant_int8"](*args, **kwargs) + + +def compress_int8_indexer_k_cache(*args, **kwargs): + return _OPS["compress_int8_indexer_k_cache"](*args, **kwargs) + + +def int8_mqa_logits(*args, **kwargs): + return _OPS["int8_mqa_logits"](*args, **kwargs) + + +def int8_paged_mqa_logits(*args, **kwargs): + return _OPS["int8_paged_mqa_logits"](*args, **kwargs) diff --git a/vllm_fl/ops/deepseek_v4_int8_indexer.py b/vllm_fl/ops/deepseek_v4_int8_indexer.py new file mode 100644 index 000000000..e6ea5a238 --- /dev/null +++ b/vllm_fl/ops/deepseek_v4_int8_indexer.py @@ -0,0 +1,922 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Direct INT8 kernels for the DeepSeek V4 sparse indexer.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _async_load(pointer, mask, other): + """Portable load used when the optional TLE-Lite package is unavailable.""" + return tl.load(pointer, mask=mask, other=other) + + +@triton.jit +def _round_to_int8(x): + """Symmetric round-half-away-from-zero, then clamp to the INT8 range. + + ``x.to(tl.int8)`` truncates toward zero, which biases every quantized + magnitude downward by up to a full LSB instead of half of one. Adding the + signed 0.5 offset before the cast recovers round-to-nearest; ``tl.clamp`` + runs afterwards so the offset itself cannot push a value out of range. + Written with plain arithmetic rather than libdevice ``rint`` to stay + portable across the vendor backends this plugin targets. + """ + return tl.clamp( + x + tl.where(x >= 0, 0.5, -0.5), -127.0, 127.0 + ).to(tl.int8) + + +@triton.jit +def fused_compress_rope_int8_indexer_cache_kernel( + state_cache, + state_stride_block, + state_stride_token, + token_to_req, + positions, + state_slots, + state_block_table, + state_block_table_stride, + state_block_size, + norm_weight, + norm_eps, + cos_sin, + cos_sin_stride, + k_cache, + kv_slots, + kv_block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + OVERLAP: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + SCALE_DIM: tl.constexpr, + KV_BLOCK_STRIDE: tl.constexpr, +): + """Compress and directly write INT8 indexer K without an FP8 stage.""" + token = tl.program_id(0) + state_slot = tl.load(state_slots + token) + if state_slot < 0: + return + position = tl.load(positions + token) + if (position + 1) % COMPRESS_RATIO != 0: + return + kv_slot = tl.load(kv_slots + token) + if kv_slot < 0: + return + + request = tl.load(token_to_req + token) + count: tl.constexpr = (1 + OVERLAP) * COMPRESS_RATIO + history = tl.arange(0, count) + history_pos = position - count + 1 + history + history_valid = history_pos >= 0 + logical_blocks = history_pos // state_block_size + physical_blocks = tl.load( + state_block_table + + request * state_block_table_stride + + logical_blocks, + mask=history_valid, + other=0, + ).to(tl.int64) + block_offsets = history_pos % state_block_size + overlap_offset = (history >= COMPRESS_RATIO).to(tl.int32) * HEAD_SIZE + dims = tl.arange(0, HEAD_SIZE) + rows = ( + state_cache + + physical_blocks * state_stride_block + + block_offsets * state_stride_token + + overlap_offset + ) + valid = history_valid[:, None] + scores = tl.load( + rows[:, None] + STATE_WIDTH + dims[None, :], + mask=valid, + other=float("-inf"), + ) + scores = tl.softmax(scores, dim=0) + values = tl.load( + rows[:, None] + dims[None, :], mask=valid, other=0.0 + ) + compressed = tl.sum(values * scores, axis=0) + + weight = tl.load(norm_weight + dims) + variance = tl.sum(compressed * compressed, axis=0) / HEAD_SIZE + normalized = compressed * tl.rsqrt(variance + norm_eps) * weight + + half_rope: tl.constexpr = ROPE_HEAD_DIM // 2 + nope_dim: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + normalized_pairs = tl.reshape(normalized, (HEAD_SIZE // 2, 2)) + normalized_even, normalized_odd = tl.split(normalized_pairs) + partner = tl.interleave(normalized_odd, normalized_even) + rope_local = dims - nope_dim + is_rope = dims >= nope_dim + pair = tl.maximum(rope_local >> 1, 0) + compressed_position = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cs = cos_sin + compressed_position * cos_sin_stride + cos_v = tl.load(cs + pair, mask=is_rope, other=1.0) + sin_v = tl.load(cs + half_rope + pair, mask=is_rope, other=0.0) + is_even = (rope_local & 1) == 0 + rotated = tl.where( + is_even, + normalized * cos_v - partner * sin_v, + normalized * cos_v + partner * sin_v, + ) + result = tl.where(is_rope, rotated, normalized).to(tl.bfloat16).to( + tl.float32 + ) + + absmax = tl.maximum(tl.max(tl.abs(result), axis=0), 1.0e-4) + scale = absmax / 127.0 + quant = _round_to_int8(result / scale) + + kv_block = kv_slot // kv_block_size + kv_pos = kv_slot % kv_block_size + block_base = kv_block.to(tl.int64) * KV_BLOCK_STRIDE + data_base = block_base + kv_pos * HEAD_SIZE + tl.store(k_cache + data_base + dims, quant) + scale_byte = block_base + kv_block_size * HEAD_SIZE + kv_pos * 4 + tl.store( + k_cache.to(tl.pointer_type(tl.float32)) + scale_byte // 4, scale + ) + + +@triton.jit +def _indexer_q_rope_int8_kernel( + positions, + q, + q_stride_t, + q_stride_h, + cos_sin, + cos_sin_stride, + q_int8, + q_int8_stride_t, + q_int8_stride_h, + weights, + weights_stride, + weights_out, + weights_out_stride, + softmax_scale, + head_scale, + HEAD_DIM: tl.constexpr, + HALF_ROPE: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + nope_dim: tl.constexpr = HEAD_DIM - 2 * HALF_ROPE + offsets = tl.arange(0, HEAD_DIM) + base = q + token * q_stride_t + head * q_stride_h + x = tl.load(base + offsets).to(tl.float32) + + position = tl.load(positions + token) + cache = cos_sin + position * cos_sin_stride + rope_local = offsets - nope_dim + is_rope = offsets >= nope_dim + partner = tl.load(base + (offsets ^ 1), mask=is_rope, other=0.0).to( + tl.float32 + ) + cs_idx = tl.maximum(rope_local >> 1, 0) + cos_v = tl.load(cache + cs_idx, mask=is_rope, other=1.0) + sin_v = tl.load(cache + HALF_ROPE + cs_idx, mask=is_rope, other=0.0) + even = (rope_local & 1) == 0 + rotated = tl.where( + even, + x * cos_v - partner * sin_v, + x * cos_v + partner * sin_v, + ) + x = tl.where(is_rope, rotated, x).to(tl.bfloat16).to(tl.float32) + + absmax = tl.maximum(tl.max(tl.abs(x), axis=0), 1.0e-4) + scale = absmax / 127.0 + quant = _round_to_int8(x / scale) + out_base = q_int8 + token * q_int8_stride_t + head * q_int8_stride_h + tl.store(out_base + offsets, quant) + + weight = tl.load(weights + token * weights_stride + head).to(tl.float32) + tl.store( + weights_out + token * weights_out_stride + head, + weight * scale * softmax_scale * head_scale, + ) + + +def fused_indexer_q_rope_quant_int8( + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply GPT-J RoPE and directly quantize indexer Q to INT8. + + The per-token/head Q scale is folded into ``weights_out``, matching the + existing FP8 indexer contract while avoiding any FP8 representation. + """ + q_int8 = torch.empty_like(q, dtype=torch.int8) + weights_out = torch.empty_like(weights, dtype=torch.float32) + _indexer_q_rope_int8_kernel[(positions.shape[0], q.shape[1])]( + positions, + q, + q.stride(0), + q.stride(1), + cos_sin_cache, + cos_sin_cache.stride(0), + q_int8, + q_int8.stride(0), + q_int8.stride(1), + weights, + weights.stride(0), + weights_out, + weights_out.stride(0), + softmax_scale, + head_scale, + HEAD_DIM=q.shape[2], + HALF_ROPE=cos_sin_cache.shape[1] // 2, + num_warps=1, + ) + return q_int8, weights_out + + +def compress_int8_indexer_k_cache( + compressor, + kv_score: torch.Tensor, + positions: torch.Tensor, + rotary_emb, +) -> None: + """Run the v0.24 compressor prologue and store indexer K directly as INT8.""" + from vllm.forward_context import get_forward_context + from vllm.models.deepseek_v4.compressor import save_partial_states + from vllm.platforms import current_platform + + kv, score = kv_score.split( + [compressor.coff * compressor.head_dim] * 2, + dim=-1, + ) + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return + + state_metadata = attn_metadata[compressor.state_cache.prefix] + token_to_req_indices = state_metadata.token_to_req_indices + slot_mapping = state_metadata.slot_mapping + num_actual = slot_mapping.shape[0] + block_table = state_metadata.block_table + block_size = state_metadata.block_size + state_cache = compressor.state_cache.kv_cache + state_width = state_cache.shape[-1] // 2 + pdl_kwargs = ( + {} + if current_platform.is_rocm() or current_platform.is_xpu() + else {"launch_pdl": False} + ) + save_partial_states( + kv=kv, + score=score, + ape=compressor.ape, + positions=positions, + state_cache=state_cache, + slot_mapping=slot_mapping, + block_size=block_size, + state_width=state_width, + compress_ratio=compressor.compress_ratio, + pdl_kwargs=pdl_kwargs, + ) + + cos_sin_cache = rotary_emb.cos_sin_cache + k_cache_metadata = attn_metadata[compressor.k_cache_prefix] + kv_cache = compressor._static_forward_context[ + compressor.k_cache_prefix + ].kv_cache + fused_compress_rope_int8_indexer_cache_kernel[(num_actual,)]( + state_cache, + state_cache.stride(0), + state_cache.stride(1), + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_table.stride(0), + block_size, + compressor.norm.weight, + compressor.rms_norm_eps, + cos_sin_cache, + cos_sin_cache.stride(0), + kv_cache, + k_cache_metadata.slot_mapping, + kv_cache.shape[1], + HEAD_SIZE=compressor.head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(compressor.head_dim), + STATE_WIDTH=state_width, + COMPRESS_RATIO=compressor.compress_ratio, + OVERLAP=compressor.overlap, + ROPE_HEAD_DIM=compressor.rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=128, + TOKEN_STRIDE=compressor.head_dim, + SCALE_DIM=4, + KV_BLOCK_STRIDE=kv_cache.stride(0), + num_warps=4, + launch_pdl=False, + ) + + +@triton.jit +def _int8_mqa_logits_h64_d128_kernel( + q, + k, + k_scale, + weights, + cu_ks, + cu_ke, + logits, + num_rows, + num_keys, + num_m_tiles, + num_n_tiles, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + PIPE_STAGES: tl.constexpr, +): + """DSV4 H64/D128 INT8 MQA with a fused query/head Tensor Core tile.""" + # Persistent scheduling: a fixed, SM-sized grid walks M tiles in a + # grid-stride loop. Each resident program keeps its Q/weight tile live + # while streaming all K tiles, instead of launching one CTA per MxN tile. + worker = tl.program_id(0) + num_workers = tl.num_programs(0) + dims = tl.arange(0, 128) + heads32 = tl.arange(0, 32) + # When M is small, spread resident CTAs across N so all SMs stay busy; + # when M is large, each CTA owns a grid-stride sequence of M tiles. A CTA + # retains Q while walking its N shard in both cases. + m_workers = tl.minimum(num_m_tiles, num_workers) + m_lane = worker % m_workers + n_lane = worker // m_workers + n_workers = tl.cdiv(num_workers, m_workers) + for pid_m in range(m_lane, num_m_tiles, m_workers): + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + row_mask = rows < num_rows + q_rows32 = tl.arange(0, BLOCK_M * 32) + query_rows = pid_m * BLOCK_M + q_rows32 // 32 + query_heads = q_rows32 % 32 + q_tile0 = tl.load( + q + query_rows[:, None] * (64 * 128) + + query_heads[:, None] * 128 + dims[None, :], + mask=(query_rows < num_rows)[:, None], other=0, + eviction_policy="evict_last", + ) + q_tile1 = tl.load( + q + query_rows[:, None] * (64 * 128) + + (query_heads[:, None] + 32) * 128 + dims[None, :], + mask=(query_rows < num_rows)[:, None], other=0, + eviction_policy="evict_last", + ) + head_weights0 = tl.load( + weights + rows[:, None] * 64 + heads32[None, :], + mask=row_mask[:, None], other=0.0, + eviction_policy="evict_last", + ) + head_weights1 = tl.load( + weights + rows[:, None] * 64 + (heads32[None, :] + 32), + mask=row_mask[:, None], other=0.0, + eviction_policy="evict_last", + ) + starts = tl.load(cu_ks + rows, mask=row_mask, other=0) + ends = tl.load(cu_ke + rows, mask=row_mask, other=0) + starts_min = tl.min(tl.where(row_mask, starts, num_keys), axis=0) + ends_max = tl.max(tl.where(row_mask, ends, 0), axis=0) + # Each packed request owns a narrow contiguous interval in the global + # K workspace. Walking all ``num_n_tiles`` made every M tile scan the + # unrelated intervals of the other requests, which becomes dominant + # for long-context/high-concurrency prefill. Start directly at this + # M tile's first useful N tile and stop after its last useful tile. + # ``n_lane`` still partitions that interval when M is too small to + # occupy all SMs by itself. + first_n_tile = starts_min // BLOCK_N + last_n_tile = tl.minimum(tl.cdiv(ends_max, BLOCK_N), num_n_tiles) + for pid_n in tl.range( + first_n_tile + n_lane, + last_n_tile, + n_workers, + num_stages=PIPE_STAGES, + ): + keys = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + key_mask = keys < num_keys + tile_min_key = pid_n * BLOCK_N + tile_max_key = tl.minimum(tile_min_key + BLOCK_N, num_keys) - 1 + # cu_ks/cu_ke describe the useful key window for every row in + # this M tile. Reject wholly out-of-window tiles before loading K + # or issuing the Tensor Core dot; the per-element mask below still + # handles partial boundary tiles. + tile_overlaps_window = ( + (tile_max_key >= starts_min) & (tile_min_key < ends_max) + ) + if tile_overlaps_window: + k_tile = _async_load( + k + keys[:, None] * 128 + dims[None, :], + key_mask[:, None], 0, + ) + scales = _async_load(k_scale + keys, key_mask, 0.0) + dots0 = tl.dot(k_tile, tl.trans(q_tile0), out_dtype=tl.int32) + scores0 = tl.reshape(dots0, (BLOCK_N, BLOCK_M, 32)).to( + tl.float32 + ) + values = tl.sum( + tl.maximum(scores0, 0.0) * head_weights0[None, :, :], + axis=2, + ) + dots1 = tl.dot(k_tile, tl.trans(q_tile1), out_dtype=tl.int32) + scores1 = tl.reshape(dots1, (BLOCK_N, BLOCK_M, 32)).to( + tl.float32 + ) + values += tl.sum( + tl.maximum(scores1, 0.0) * head_weights1[None, :, :], + axis=2, + ) + values *= scales[:, None] + valid = (row_mask[:, None] & key_mask[None, :] + & (keys[None, :] >= starts[:, None]) + & (keys[None, :] < ends[:, None])) + tl.store( + logits + rows[:, None] * num_keys + keys[None, :], + tl.trans(values), mask=valid, + eviction_policy="evict_first", + ) + + +@triton.jit +def _int8_mqa_logits_tensor_core_kernel( + q, + q_stride_m, + q_stride_h, + q_stride_d, + k, + k_stride_n, + k_stride_d, + k_scale, + weights, + weights_stride_m, + weights_stride_h, + cu_ks, + cu_ke, + logits, + logits_stride_m, + num_rows, + num_keys, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Stride-aware parameterized INT8 Tensor Core MQA fallback.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + keys = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = rows < num_rows + key_mask = keys < num_keys + dims = tl.arange(0, BLOCK_D) + dim_mask = dims < HEAD_DIM + + # Pad heads/dims to Tensor Core-compatible powers of two. Padded query + # rows and weights are zero, so they disappear in the head reduction. + q_rows = tl.arange(0, BLOCK_M * BLOCK_H) + query_rows = pid_m * BLOCK_M + q_rows // BLOCK_H + query_heads = q_rows % BLOCK_H + q_tile = tl.load( + q + + query_rows[:, None] * q_stride_m + + query_heads[:, None] * q_stride_h + + dims[None, :] * q_stride_d, + mask=(query_rows < num_rows)[:, None] + & (query_heads < NUM_HEADS)[:, None] + & dim_mask[None, :], + other=0, + eviction_policy="evict_last", + ) + k_tile = tl.load( + k + keys[:, None] * k_stride_n + dims[None, :] * k_stride_d, + mask=key_mask[:, None] & dim_mask[None, :], + other=0, + eviction_policy="evict_first", + ) + scales = tl.load(k_scale + keys, mask=key_mask, other=0.0) + + dots = tl.dot(k_tile, tl.trans(q_tile), out_dtype=tl.int32) + scores = tl.reshape(dots, (BLOCK_N, BLOCK_M, BLOCK_H)).to( + tl.float32 + ) + heads = tl.arange(0, BLOCK_H) + head_weights = tl.load( + weights + + rows[:, None] * weights_stride_m + + heads[None, :] * weights_stride_h, + mask=row_mask[:, None] & (heads < NUM_HEADS)[None, :], + other=0.0, + eviction_policy="evict_last", + ) + values = tl.sum( + tl.maximum(scores * scales[:, None, None], 0.0) + * head_weights[None, :, :], + axis=2, + ) + + starts = tl.load(cu_ks + rows, mask=row_mask, other=0) + ends = tl.load(cu_ke + rows, mask=row_mask, other=0) + valid = ( + row_mask[:, None] + & key_mask[None, :] + & (keys[None, :] >= starts[:, None]) + & (keys[None, :] < ends[:, None]) + ) + tl.store( + logits + + rows[:, None] * logits_stride_m + + keys[None, :], + tl.trans(values), + mask=valid, + eviction_policy="evict_first", + ) + + +def int8_mqa_logits( + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_ks: torch.Tensor, + cu_ke: torch.Tensor, +) -> torch.Tensor: + """Contiguous INT8 indexer logits used by prefill.""" + num_keys = k.shape[0] + logits = torch.empty( + (q.shape[0], num_keys), dtype=torch.float32, device=q.device + ) + if ( + q.shape[1] == 64 + and q.shape[2] == 128 + and q.is_contiguous() + and k.is_contiguous() + and weights.is_contiguous() + ): + # A 128-key tile needs enough registers/shared memory that Hopper can + # keep only one 8-warp CTA resident per SM. Halving N preserves the + # K/Q traffic while allowing two resident CTAs; the larger worker grid + # then hides the long INT8 MMA/reduction latency on long prefills. + block_m, block_n = 4, 64 + num_m_tiles = triton.cdiv(q.shape[0], block_m) + num_n_tiles = triton.cdiv(num_keys, block_n) + sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count + total_tiles = num_m_tiles * num_n_tiles + target_workers = min(2 * sm_count, total_tiles) + # Every M lane must own the same number of N shards. A partially + # filled final group would make ``ceil(num_workers / m_workers)`` skip + # one shard for the M lanes that do not have a corresponding CTA. + m_workers = min(num_m_tiles, target_workers) + n_workers = min(max(target_workers // m_workers, 1), num_n_tiles) + num_workers = m_workers * n_workers + _int8_mqa_logits_h64_d128_kernel[(num_workers,)]( + q, + k, + k_scale, + weights, + cu_ks, + cu_ke, + logits, + q.shape[0], + num_keys, + num_m_tiles, + num_n_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + PIPE_STAGES=1, + num_warps=8, + num_stages=1, + ) + return logits + + block_h = triton.next_power_of_2(q.shape[1]) + block_d = max(32, triton.next_power_of_2(q.shape[2])) + # Match DeepGEMM's 128 query-head rows per MMA tile where possible. + block_m = max(1, 128 // block_h) + block_n = 64 + _int8_mqa_logits_tensor_core_kernel[ + (triton.cdiv(q.shape[0], block_m), triton.cdiv(num_keys, block_n)) + ]( + q, + q.stride(0), + q.stride(1), + q.stride(2), + k, + k.stride(0), + k.stride(1), + k_scale, + weights, + weights.stride(0), + weights.stride(1), + cu_ks, + cu_ke, + logits, + logits.stride(0), + q.shape[0], + num_keys, + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_H=block_h, + BLOCK_D=block_d, + num_warps=8, + num_stages=1, + ) + return logits + + +@triton.jit +def _int8_paged_mqa_logits_h64_d128_kernel( + q, + cache, + cache_block_stride, + weights, + context_lens, + context_row_stride, + context_col_stride, + block_table, + block_table_stride, + logits, + max_model_len, + NEXT_N: tl.constexpr, + PAGES_PER_CTA: tl.constexpr, +): + """DSV4-specialized H=64, D=128, page=64 decode kernel. + + This mirrors FlagGems' BF16 specialization: source-level shape literals, + a small fixed group of physical pages per program, and cache hints matching + the reuse distance of Q/K. Grouping amortizes Q/weight loads without + collapsing the two-dimensional page parallelism. + """ + row = tl.program_id(0) + first_logical_block = tl.program_id(1) * PAGES_PER_CTA + batch = row // NEXT_N + next_idx = row % NEXT_N + # Strides come from the caller: a compact (B, 1) ``context_lens`` passes a + # zero column stride so the shared length broadcasts across next_n. + context_len = tl.load( + context_lens + + batch * context_row_stride + + next_idx * context_col_stride + ) + first_key_start = first_logical_block * 64 + if first_key_start >= context_len: + return + + heads = tl.arange(0, 64) + dims = tl.arange(0, 128) + positions = tl.arange(0, 64) + + # Decode Q and weights are contiguous and reused by every cache page. + q_tile = tl.load( + q + row * (64 * 128) + heads[:, None] * 128 + dims[None, :], + eviction_policy="evict_last", + ) + head_weights = tl.load( + weights + row * 64 + heads, eviction_policy="evict_last" + ) + + for page_offset in tl.static_range(0, PAGES_PER_CTA): + logical_block = first_logical_block + page_offset + key_start = logical_block * 64 + page_valid = key_start < context_len + physical_block = tl.load( + block_table + batch * block_table_stride + logical_block, + mask=page_valid, + other=0, + ).to(tl.int64) + page_base = physical_block * cache_block_stride + k_tile = tl.load( + cache + page_base + positions[:, None] * 128 + dims[None, :], + mask=page_valid, + other=0, + eviction_policy="evict_first", + ).to(tl.int8) + k_scales = tl.load( + cache.to(tl.pointer_type(tl.float32)) + + (page_base + 64 * 128) // 4 + + positions, + mask=page_valid, + other=0.0, + eviction_policy="evict_first", + ) + + dots = tl.dot(k_tile, tl.trans(q_tile), out_dtype=tl.int32) + activated = tl.maximum( + dots.to(tl.float32) * k_scales[:, None], 0.0 + ) + result = tl.sum(activated * head_weights[None, :], axis=1) + output = logits + row * max_model_len + key_start + positions + tl.store( + output, + result, + mask=page_valid & (positions < context_len - key_start), + eviction_policy="evict_first", + ) + + +@triton.jit +def _int8_paged_mqa_logits_kernel( + q, + q_stride_b, + q_stride_n, + q_stride_h, + cache, + cache_block_stride, + block_size, + weights, + weights_stride, + context_lens, + context_row_stride, + context_col_stride, + block_table, + block_table_stride, + logits, + logits_stride, + max_model_len, + NEXT_N: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_KEYS: tl.constexpr, + BLOCK_HEADS: tl.constexpr, +): + row = tl.program_id(0) + key_block = tl.program_id(1) + batch = row // NEXT_N + next_idx = row % NEXT_N + # See int8_paged_mqa_logits: a compact (B, 1) ``context_lens`` passes a zero + # column stride so the shared length broadcasts across next_n positions. + context_len = tl.load( + context_lens + + batch * context_row_stride + + next_idx * context_col_stride + ) + + # The logits buffer is deliberately not cleaned: the downstream top-k + # kernel is length-masked. Decode allocates it at max_model_len, which can + # be tens of thousands for DSV4, while a request commonly has only a few + # cache blocks. Returning before any page-table/cache access is critical; + # otherwise every padded tile performs a complete 64-head dot product. + key_start = key_block * BLOCK_KEYS + if key_start >= context_len: + return + keys = key_start + tl.arange(0, BLOCK_KEYS) + key_mask = (keys < context_len) & (keys < max_model_len) + logical_blocks = keys // block_size + block_offsets = keys % block_size + physical_blocks = tl.load( + block_table + batch * block_table_stride + logical_blocks, + mask=key_mask, + other=0, + ).to(tl.int64) + byte_bases = physical_blocks * cache_block_stride + block_offsets * HEAD_DIM + dims = tl.arange(0, HEAD_DIM) + k_tile = tl.load( + cache + byte_bases[:, None] + dims[None, :], + mask=key_mask[:, None], + other=0, + ).to(tl.int8) + scale_offsets = ( + physical_blocks * cache_block_stride + + block_size * HEAD_DIM + + block_offsets * 4 + ) + k_scales = tl.load( + cache.to(tl.pointer_type(tl.float32)) + scale_offsets // 4, + mask=key_mask, + other=0.0, + ) + + acc = tl.zeros((BLOCK_KEYS,), tl.float32) + # Map INT8 K@Q to Hopper tensor cores. The previous scalar-head loop cast + # both operands to FP32 and emitted 64 independent reductions per key + # tile. Tiling heads makes the contraction a native INT8 dot producing + # exact INT32 accumulators; scales/ReLU/weights remain FP32 as required by + # the indexer definition. + for head_start in tl.static_range(0, NUM_HEADS, BLOCK_HEADS): + heads = head_start + tl.arange(0, BLOCK_HEADS) + head_mask = heads < NUM_HEADS + q_tile = tl.load( + q + + batch * q_stride_b + + next_idx * q_stride_n + + heads[:, None] * q_stride_h + + dims[None, :], + mask=head_mask[:, None], + other=0, + ) + dots = tl.dot(k_tile, tl.trans(q_tile), out_dtype=tl.int32) + head_weights = tl.load( + weights + row * weights_stride + heads, + mask=head_mask, + other=0.0, + ) + activated = tl.maximum( + dots.to(tl.float32) * k_scales[:, None], 0.0 + ) + acc += tl.sum(activated * head_weights[None, :], axis=1) + tl.store(logits + row * logits_stride + keys, acc, mask=key_mask) + + +def int8_paged_mqa_logits( + q: torch.Tensor, + cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + max_model_len: int, +) -> torch.Tensor: + """Paged INT8 indexer logits used by decode.""" + batch, next_n, num_heads, head_dim = q.shape + logits = torch.empty( + (batch * next_n, max_model_len), dtype=torch.float32, device=q.device + ) + flat = cache.view(torch.uint8) + + # ``context_lens`` is (B, next_n) under native spec decode but (B, 1) + # otherwise. The kernels index it as ``batch * row_stride + next_idx``, so a + # compact single-column tensor must advertise a zero column stride to + # broadcast the shared length across the next_n speculative positions -- + # taking stride(0) alone would walk into the following request's entry. + context_row_stride = context_lens.stride(0) + if context_lens.dim() > 1 and context_lens.shape[1] == 1: + context_col_stride = 0 + elif context_lens.dim() == 1: + context_row_stride, context_col_stride = context_lens.stride(0), 0 + else: + context_col_stride = context_lens.stride(1) + if ( + num_heads == 64 + and head_dim == 128 + and cache.shape[1] == 64 + and q.is_contiguous() + and weights.is_contiguous() + and context_lens.is_contiguous() + ): + # A context cannot address more pages than its block-table row. Avoid + # launching capacity-only CTAs when metadata carries a compact table. + grid_blocks = min( + triton.cdiv(max_model_len, 64), block_table.shape[1] + ) + pages_per_cta = 2 + _int8_paged_mqa_logits_h64_d128_kernel[ + (batch * next_n, triton.cdiv(grid_blocks, pages_per_cta)) + ]( + q, + flat, + cache.stride(0), + weights, + context_lens, + context_row_stride, + context_col_stride, + block_table, + block_table.stride(0), + logits, + max_model_len, + NEXT_N=next_n, + PAGES_PER_CTA=pages_per_cta, + num_warps=4, + num_stages=1, + ) + return logits + + _int8_paged_mqa_logits_kernel[ + (batch * next_n, triton.cdiv(max_model_len, 128)) + ]( + q, + q.stride(0), + q.stride(1), + q.stride(2), + flat, + cache.stride(0), + cache.shape[1], + weights, + weights.stride(0), + context_lens, + context_row_stride, + context_col_stride, + block_table, + block_table.stride(0), + logits, + logits.stride(0), + max_model_len, + NEXT_N=next_n, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + BLOCK_KEYS=128, + BLOCK_HEADS=32, + num_warps=8, + ) + return logits diff --git a/vllm_fl/ops/deepseek_v4_int8_woa.py b/vllm_fl/ops/deepseek_v4_int8_woa.py new file mode 100644 index 000000000..af0449a54 --- /dev/null +++ b/vllm_fl/ops/deepseek_v4_int8_woa.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused inverse-RoPE and INT8 activation quantization for DSV4 wo_a.""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + +from vllm_fl.dispatch import resolve_op + +DSV4_INV_ROPE_QUANT_INT8_OP = "deepseek_v4_inv_rope_quant_int8" +# The output projection runs inside the full-graph compiled model. Resolve the +# backend before tracing so Dynamo never enters OpManager's synchronization +# path; backend selection still goes through OpManager once at module import. +_dispatch_inv_rope_quant_int8 = resolve_op(DSV4_INV_ROPE_QUANT_INT8_OP) + + +@triton.jit +def _inv_rope_quant_int8_kernel( + o, + positions, + cos_sin, + o_q, + o_scale, + o_stride_t, + o_stride_h, + cs_stride_t, + q_stride_t, + q_stride_g, + scale_stride_t, + scale_stride_g, + NUM_TOKENS, + HEADS_PER_GROUP: tl.constexpr, + HEAD_DIM: tl.constexpr, + NOPE_DIM: tl.constexpr, + HALF_ROPE: tl.constexpr, + GROUP_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + group = tl.program_id(1).to(tl.int64) + offsets = tl.arange(0, BLOCK_D) + valid = offsets < GROUP_DIM + head = offsets // HEAD_DIM + dim = offsets % HEAD_DIM + source_head = group * HEADS_PER_GROUP + head + ptr = o + token * o_stride_t + source_head * o_stride_h + dim + x = tl.load(ptr, mask=valid, other=0.0).to(tl.float32) + + rope_local = dim - NOPE_DIM + is_rope = valid & (dim >= NOPE_DIM) + partner = tl.load( + ptr + tl.where((rope_local & 1) == 0, 1, -1), + mask=is_rope, + other=0.0, + ).to(tl.float32) + pos = tl.load(positions + token, mask=token < NUM_TOKENS, other=0) + cs = cos_sin + pos * cs_stride_t + pair = tl.maximum(rope_local >> 1, 0) + cos_v = tl.load(cs + pair, mask=is_rope, other=1.0) + sin_v = tl.load(cs + HALF_ROPE + pair, mask=is_rope, other=0.0) + inv = tl.where( + (rope_local & 1) == 0, + x * cos_v + partner * sin_v, + x * cos_v - partner * sin_v, + ) + x = tl.where(is_rope, inv, x).to(tl.bfloat16).to(tl.float32) + + absmax = tl.maximum(tl.max(tl.where(valid, tl.abs(x), 0.0), axis=0), 1.0e-4) + scale = absmax / 127.0 + rounded = x / scale + tl.where(x >= 0, 0.5, -0.5) + quant = tl.clamp(rounded, -127.0, 127.0).to(tl.int8) + tl.store( + o_q + token * q_stride_t + group * q_stride_g + offsets, + quant, + mask=valid, + ) + tl.store( + o_scale + token * scale_stride_t + group * scale_stride_g, + scale, + ) + + +def fused_inv_rope_quant_int8_triton( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return group-major symmetric INT8 activations and row scales.""" + tokens, heads, head_dim = o.shape + assert heads == n_groups * heads_per_group + assert head_dim == nope_dim + rope_dim + group_dim = heads_per_group * head_dim + block_d = triton.next_power_of_2(group_dim) + o_q = torch.empty( + (n_groups, tokens, group_dim), + dtype=torch.int8, + device=o.device, + ) + o_scale = torch.empty( + (n_groups, tokens, 1), + dtype=torch.float32, + device=o.device, + ) + _inv_rope_quant_int8_kernel[(tokens, n_groups)]( + o, + positions, + cos_sin_cache, + o_q, + o_scale, + o.stride(0), + o.stride(1), + cos_sin_cache.stride(0), + o_q.stride(1), + o_q.stride(0), + o_scale.stride(1), + o_scale.stride(0), + tokens, + HEADS_PER_GROUP=heads_per_group, + HEAD_DIM=head_dim, + NOPE_DIM=nope_dim, + HALF_ROPE=rope_dim // 2, + GROUP_DIM=group_dim, + BLOCK_D=block_d, + num_warps=8, + num_stages=3, + ) + return o_q, o_scale + + +def fused_inv_rope_quant_int8( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Dispatch DSV4 inverse-RoPE INT8 quantization through OpManager.""" + return _dispatch_inv_rope_quant_int8( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) diff --git a/vllm_fl/ops/sparse_attn_indexer.py b/vllm_fl/ops/sparse_attn_indexer.py new file mode 100644 index 000000000..75b474611 --- /dev/null +++ b/vllm_fl/ops/sparse_attn_indexer.py @@ -0,0 +1,535 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom Sparse Attention Indexer layers.""" + +import torch + +import vllm.envs as envs +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, +) +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, +) +from vllm.v1.worker.workspace import current_workspace_manager + +from vllm_fl.dispatch import resolve_op +from vllm_fl.ops.deepseek_v4 import ( + int8_mqa_logits, + int8_paged_mqa_logits, +) + +# Resolve before torch.compile traces sparse_attn_indexer_fl. Lazy CachedOp +# lookup would make Dynamo enter OpManager's synchronization path. Backend +# selection still belongs to OpManager; the compiled graph only sees the +# selected implementation. +_indexer_k_quant_and_cache = resolve_op("indexer_k_quant_and_cache") +_cp_gather_indexer_k_quant_cache = resolve_op("cp_gather_indexer_k_quant_cache") +_top_k_per_row_prefill = resolve_op("top_k_per_row_prefill") +_pack_seq_triton = resolve_op("pack_seq_triton") +_top_k_per_row_decode = resolve_op("top_k_per_row_decode") +_unpack_seq_triton = resolve_op("unpack_seq_triton") + +logger = init_logger(__name__) + +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 + +# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. +MXFP4_BLOCK_SIZE = 32 + + +def _gather_workspace_shapes( + total_seq_lens: int, + head_dim: int, + fp8_dtype: torch.dtype, + use_fp4_cache: bool, + use_int8_cache: bool = False, +) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: + """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for + the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 + scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" + if use_fp4_cache: + return ( + ((total_seq_lens, head_dim // 2), torch.uint8), + ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), + ) + if use_int8_cache: + return ( + ((total_seq_lens, head_dim), torch.int8), + ((total_seq_lens, 4), torch.uint8), + ) + return ( + ((total_seq_lens, head_dim), fp8_dtype), + ((total_seq_lens, 4), torch.uint8), + ) + + +def kv_cache_as_quant_view( + kv_cache: torch.Tensor, + head_dim: int, + use_fp4_cache: bool, +) -> torch.Tensor: + """4D ``[num_blocks, block_size, 1, head_width]`` view expected by + DeepGEMM, from the 3D indexer kv-cache allocation.""" + if use_fp4_cache: + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 + num_blocks, block_size, _ = kv_cache.shape + page_bytes = int(kv_cache.stride(0)) + fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE + return torch.as_strided( + kv_cache, + size=(num_blocks, block_size, 1, fp4_bytes), + stride=(page_bytes, fp4_bytes, fp4_bytes, 1), + ) + return kv_cache.unsqueeze(-2) + + +def sparse_attn_indexer_fl( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor, + skip_k_cache_insert: bool, + use_fp4_cache: bool = False, +) -> torch.Tensor: + # careful! this will be None in dummy run + attn_metadata = get_forward_context().attn_metadata + fp8_dtype = current_platform.fp8_dtype() + k_cache_prefix = _resolve_layer_name(k_cache_prefix) + use_int8_cache = q_quant.dtype == torch.int8 + + # assert isinstance(attn_metadata, dict) + if not isinstance(attn_metadata, dict): + # Reserve workspace for indexer during profiling run + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, + head_dim, + fp8_dtype, + use_fp4_cache, + use_int8_cache, + ) + current_workspace_manager().get_simultaneous( + values_spec, + scales_spec, + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + ) + + # Dummy allocation to simulate for peak logits tensor memory during inference. + # FP8 elements so elements == bytes + max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + _ = torch.empty( + max_logits_elems, dtype=torch.uint8, device=hidden_states.device + ) + + return sparse_attn_indexer_fl_fake( + hidden_states, + k_cache_prefix, + kv_cache, + q_quant, + q_scale, + k, + weights, + quant_block_size, + scale_fmt, + topk_tokens, + head_dim, + max_model_len, + total_seq_lens, + topk_indices_buffer, + skip_k_cache_insert, + use_fp4_cache, + ) + attn_metadata_narrowed = attn_metadata[k_cache_prefix] + assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) + slot_mapping = attn_metadata_narrowed.slot_mapping + has_decode = attn_metadata_narrowed.num_decodes > 0 + has_prefill = attn_metadata_narrowed.num_prefills > 0 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + + # q_scale is required iff the FP4 cache path is enabled; the FP8 path + # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. + if use_fp4_cache: + assert q_scale is not None, "use_fp4_cache=True requires q_scale" + else: + assert q_scale is None, "q_scale must be None when use_fp4_cache=False" + + # During speculative decoding, k may be padded to the CUDA graph batch + # size while slot_mapping only covers actual tokens. Truncate k to avoid + # out-of-bounds reads in the kernel. + num_tokens = slot_mapping.shape[0] + if k is not None: + k = k[:num_tokens] + + if not skip_k_cache_insert: + # scale_fmt can be None, but the function expects str + assert scale_fmt is not None + assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" + _indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + ) + + topk_indices_buffer[: hidden_states.shape[0]] = -1 + if has_prefill: + prefill_metadata = attn_metadata_narrowed.prefill + assert prefill_metadata is not None + + # Get the full shared workspace buffers once (will allocate on first use). + # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and + # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 + # scales) based on use_fp4_cache. + workspace_manager = current_workspace_manager() + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, + head_dim, + fp8_dtype, + use_fp4_cache, + use_int8_cache, + ) + k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + values_spec, + scales_spec, + ) + for chunk in prefill_metadata.chunks: + k_quant = k_quant_full[: chunk.total_seq_lens] + k_scale = k_scale_full[: chunk.total_seq_lens] + + if not chunk.skip_kv_gather: + _cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.cu_seq_lens, + ) + + q_slice = q_quant[chunk.token_start : chunk.token_end] + q_scale_slice = ( + q_scale[chunk.token_start : chunk.token_end] + if q_scale is not None + else None + ) + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. + if use_fp4_cache: + q_slice_cast = q_slice.view(torch.int8) + k_quant_cast = k_quant.view(torch.int8) + k_scale_cast = k_scale.view(torch.int32).squeeze(-1) + else: + q_slice_cast = q_slice + k_quant_cast = k_quant + k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + + if use_int8_cache: + logits = int8_mqa_logits( + q_slice_cast, + k_quant_cast, + k_scale_cast, + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + ) + else: + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + weights[chunk.token_start : chunk.token_end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + clean_logits=False, + ) + num_rows = logits.shape[0] + + topk_indices = topk_indices_buffer[ + chunk.token_start : chunk.token_end, :topk_tokens + ] + + _top_k_per_row_prefill( + logits, + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + topk_indices, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + + if has_decode: + decode_metadata = attn_metadata_narrowed.decode + assert decode_metadata is not None + if not use_int8_cache: + kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) + decode_lens = decode_metadata.decode_lens + if decode_metadata.requires_padding: + # pad in edge case where we have short chunked prefill length < + # decode_threshold since we unstrictly split + # prefill and decode by decode_threshold + # (currently set to 1 + speculative tokens). + # FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK — + # downstream context_lens masks stale slots). MXFP4 Q is two + # uint8 tensors (values + ue8m0 scales) — use the dedicated uint8 + # packer with pad_byte=0 so padded slots dequantize to 0 and + # can't produce NaN/Inf in the logits kernel. + if q_scale is not None: + padded_q_quant_decode_tokens = _pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens, pad_value=0 + ) + padded_q_scale = _pack_seq_triton( + q_scale[:num_decode_tokens], decode_lens, pad_value=0 + ) + else: + padded_q_quant_decode_tokens = _pack_seq_triton( + q_quant[:num_decode_tokens], decode_lens + ) + padded_q_scale = None + else: + padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_quant.shape[1:] + ) + if q_scale is not None: + padded_q_scale = q_scale[:num_decode_tokens].reshape( + decode_lens.shape[0], -1, *q_scale.shape[1:] + ) + else: + padded_q_scale = None + # TODO: move and optimize below logic with triton kernels + batch_size = padded_q_quant_decode_tokens.shape[0] + next_n = padded_q_quant_decode_tokens.shape[1] + num_padded_tokens = batch_size * next_n + seq_lens = decode_metadata.seq_lens[:batch_size] + # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) + # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; + # the downstream topk kernels accept both 1D and 2D. + padded_q_quant_cast = ( + padded_q_quant_decode_tokens.view(torch.int8) + if use_fp4_cache + else padded_q_quant_decode_tokens + ) + + if use_int8_cache: + logits = int8_paged_mqa_logits( + padded_q_quant_decode_tokens, + kv_cache, + weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + max_model_len=max_model_len, + ) + else: + logits = fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + clean_logits=False, + ) + num_rows = logits.shape[0] + topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] + _top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + + # if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048): + # workspace_manager = current_workspace_manager() + # (topk_workspace,) = workspace_manager.get_simultaneous( + # ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), + # ) + # torch.ops._C.persistent_topk( + # logits, + # seq_lens, + # topk_indices, + # topk_workspace, + # topk_tokens, + # attn_metadata_narrowed.max_seq_len, + # ) + # else: + # torch.ops._C.top_k_per_row_decode( + # logits, + # next_n, + # seq_lens, + # topk_indices, + # num_rows, + # logits.stride(0), + # logits.stride(1), + # topk_tokens, + # ) + + if decode_metadata.requires_padding: + # if padded, we need to unpack + # the topk indices removing padded tokens + topk_indices = _unpack_seq_triton( + topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]), + decode_lens, + ) + topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = ( + topk_indices + ) + + return topk_indices_buffer + + +def sparse_attn_indexer_fl_fake( + hidden_states: torch.Tensor, + k_cache_prefix: LayerNameType, + kv_cache: torch.Tensor, + q_quant: torch.Tensor, + q_scale: torch.Tensor | None, + k: torch.Tensor, + weights: torch.Tensor, + quant_block_size: int, + scale_fmt: str | None, + topk_tokens: int, + head_dim: int, + max_model_len: int, + total_seq_lens: int, + topk_indices_buffer: torch.Tensor | None, + skip_k_cache_insert: bool, + use_fp4_cache: bool = False, +) -> torch.Tensor: + return topk_indices_buffer + + +direct_register_custom_op( + op_name="sparse_attn_indexer_fl", + op_func=sparse_attn_indexer_fl, + mutates_args=["topk_indices_buffer"], + fake_impl=sparse_attn_indexer_fl_fake, + dispatch_key=current_platform.dispatch_key, +) + + +class SparseAttnIndexerFL(SparseAttnIndexer): + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a + separate custom op since it involves heavy custom kernels like `mqa_logits`, + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires + specific memory layout or implementation for different hardware backends to + achieve optimal performance. + + For now, the default native path will use CUDA backend path. Other platform + may requires add the corresponding Custom Op name `sparse_attn_indexer` to + `custom_ops` in `CompilationConfig` to enable the platform specific path. + """ + + # def __init__( + # self, + # k_cache, + # quant_block_size: int, + # scale_fmt: str, + # topk_tokens: int, + # head_dim: int, + # max_model_len: int, + # max_total_seq_len: int, + # topk_indices_buffer: torch.Tensor, + # skip_k_cache_insert: bool = False, + # use_fp4_cache: bool = False, + # ): + # super().__init__() + # self.k_cache = k_cache + # self.quant_block_size = quant_block_size + # self.scale_fmt = scale_fmt + # self.topk_tokens = topk_tokens + # self.head_dim = head_dim + # self.max_model_len = max_model_len + # self.max_total_seq_len = max_total_seq_len + # self.topk_indices_buffer = topk_indices_buffer + # self.skip_k_cache_insert = skip_k_cache_insert + # self.use_fp4_cache = use_fp4_cache + # if current_platform.is_cuda() and not has_deep_gemm(): + # raise RuntimeError( + # "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." + # ) + + def forward_native( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + ): + return self.forward_oot(hidden_states, q_quant, k, weights) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + ): + """Keep NVIDIA CUDA on the FL INT8-aware indexer path. + + vLLM 0.24 selects its in-tree CUDA platform even when the FL plugin is + loaded. Without this override ``CustomOp.forward`` resolves the + inherited upstream ``SparseAttnIndexer.forward_cuda``, which sends an + INT8 query to DeepGEMM's FP8-only MQA kernel. The FL entrypoint owns + the mixed-cache dispatch and routes its helper kernels through + OpManager. + """ + return self.forward_oot(hidden_states, q_quant, k, weights) + + def forward_oot( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + weights: torch.Tensor, + ): + # FP8 path: single tensor (per-token scale is folded into `weights`). + # FP4 path: (values, scales) tuple with scales required by the kernel. + if isinstance(q_quant, tuple): + q_values, q_scale = q_quant + else: + q_values, q_scale = q_quant, None + return torch.ops.vllm.sparse_attn_indexer_fl( + hidden_states, + _encode_layer_name(self.k_cache.prefix), + self.k_cache.kv_cache, + q_values, + q_scale, + k, + weights, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + self.skip_k_cache_insert, + self.use_fp4_cache, + ) diff --git a/vllm_fl/quantization/w8a8/linear.py b/vllm_fl/quantization/w8a8/linear.py index 774c99c33..c9e350816 100644 --- a/vllm_fl/quantization/w8a8/linear.py +++ b/vllm_fl/quantization/w8a8/linear.py @@ -93,26 +93,78 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: setattr(layer, i_zp_name, None) setattr(layer, azp_adj_name, None) + # DSV4 wo_a is logically a grouped BMM. Keep a group-major copy in + # CUTLASS column-major layout: selecting a contiguous [N, K] group + # and transposing it yields a zero-copy [K, N] view with stride + # (1, K). This makes the graph path use vLLM's registered CUDA op. + if getattr(layer, "is_bmm", False): + num_groups = int(getattr(layer, "bmm_batch_size", 0)) + prepared_weight = getattr(layer, w_q_name) + prepared_scale = getattr(layer, w_s_name) + if num_groups <= 0: + raise ValueError( + "FL W8A8 grouped linear requires bmm_batch_size" + ) + if prepared_weight.shape[1] % num_groups: + raise ValueError( + "FL W8A8 grouped output size must be divisible by " + "bmm_batch_size" + ) + output_per_group = prepared_weight.shape[1] // num_groups + grouped_weight = prepared_weight.t().view( + num_groups, + output_per_group, + prepared_weight.shape[0], + ).contiguous() + grouped_scale = prepared_scale.view( + num_groups, output_per_group + ).contiguous() + layer.register_buffer( + "_fl_w8a8_grouped_weight", grouped_weight, persistent=False + ) + layer.register_buffer( + "_fl_w8a8_grouped_weight_scale", + grouped_scale, + persistent=False, + ) + def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - from flag_gems import scaled_mm - weight, weight_scale, _, _, _ = self._get_layer_params(layer) original_shape = x.shape x_2d = x.reshape(-1, original_shape[-1]).contiguous() - x_q, x_scale = _dynamic_per_token_quant_int8(x_2d) - output = scaled_mm( - x_q, - weight, - x_scale, - weight_scale, - bias=bias, - out_dtype=x.dtype, - ) + if torch.compiler.is_compiling(): + # Keep Dynamo away from the runtime policy manager and expose + # graph-safe registered CUDA ops to Inductor/CUDA Graph. + from vllm import _custom_ops as ops + + x_q, x_scale, _ = ops.scaled_int8_quant( + x_2d, None, None, symmetric=True + ) + output = ops.cutlass_scaled_mm( + x_q, + weight, + scale_a=x_scale, + scale_b=weight_scale, + bias=bias, + out_dtype=x.dtype, + ) + else: + from flag_gems import scaled_mm + + x_q, x_scale = _dynamic_per_token_quant_int8(x_2d) + output = scaled_mm( + x_q, + weight, + x_scale, + weight_scale, + bias=bias, + out_dtype=x.dtype, + ) return output.reshape(*original_shape[:-1], weight.shape[1])