diff --git a/.lil/changes/vllm-dense-mla-partial-precision.json b/.lil/changes/vllm-dense-mla-partial-precision.json new file mode 100644 index 000000000000..e2dd270a6214 --- /dev/null +++ b/.lil/changes/vllm-dense-mla-partial-precision.json @@ -0,0 +1,15 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-dense-mla-partial-precision", + "category": "feature", + "summary": "Select BF16 or FP32 split partials for B12X dense MLA", + "models": ["Kimi-K3"], + "compatibility": "BF16 MLA partials remain the default. VLLM_K3_DENSE_MLA_PARTIAL_DTYPE=fp32 requires the B12X FP32-partials companion.", + "details": [ + "Passes the declared partial dtype to prepared B12X MLA plans.", + "Covers both partial dtypes in CUDA Graph replay with capacity changes and page addresses beyond the signed 32-bit offset boundary." + ], + "pull_requests": [], + "authors": ["voipmonitor"], + "requires": ["b12x-dense-mla-fp32-partials"] +} diff --git a/.lil/changes/vllm-kimi-b12x-kda-prefill.json b/.lil/changes/vllm-kimi-b12x-kda-prefill.json new file mode 100644 index 000000000000..69589a449c8d --- /dev/null +++ b/.lil/changes/vllm-kimi-b12x-kda-prefill.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-kimi-b12x-kda-prefill", + "category": "feature", + "summary": "Run Kimi-K3 KDA prefill directly against the recurrent state pool", + "models": [ + "Kimi-K3" + ], + "compatibility": "Select --kda-prefill-backend b12x; auto backend selection is unchanged.", + "details": [ + "Prepares caller-owned scratch before execution and avoids gathering and scattering recurrent states." + ], + "pull_requests": [], + "requires": [], + "authors": [ + "voipmonitor" + ] +} diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 67282a095005..4824ffad899d 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -29,6 +29,7 @@ _flashinfer_kda_prefill, _flashkda_prefill, _store_cache_checkpoints_kernel, + is_b12x_kda_prefill_supported, is_flashinfer_fused_kda_decode_supported, is_flashinfer_recurrent_kda_prefill_supported, is_flashkda_supported, @@ -50,6 +51,7 @@ ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.utils.b12x import get_b12x_kda_prefill from vllm.utils.flashinfer import flashinfer_fused_kda_decode from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -1366,6 +1368,10 @@ def _require_kda_prefill_backend( supported = is_flashinfer_recurrent_kda_prefill_supported( 128, torch.bfloat16, state_dtype, lower_bound ) + elif backend == "b12x": + supported = is_b12x_kda_prefill_supported( + 128, torch.bfloat16, state_dtype, lower_bound + ) else: assert backend == "flashkda" supported = is_flashkda_supported(128, torch.bfloat16, state_dtype, lower_bound) @@ -1411,6 +1417,23 @@ def _run_kda_prefill_backend( seq_order=seq_order, ) + if backend == "b12x": + return _b12x_kda_prefill( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + has_initial_state=torch.ones( + initial_state.shape[0], dtype=torch.bool, device=q.device + ), + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + assert backend == "flashkda" import vllm._flashkda_C # noqa: F401 @@ -1439,6 +1462,110 @@ def _run_kda_prefill_backend( ) +def _b12x_kda_prefill( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + lower_bound: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run b12x prefill with the layer's slot conventions. + + Slot ``NULL_BLOCK_ID`` is the null slot. Request ``r`` owns slot ``r + 1`` + of a state pool, reads it only when ``has_initial_state[r]`` is set, and + receives its final state there; no checkpoint slot is written. + """ + from b12x.preparation import PreparationSession, PreparedCall, require_prepared + + api = get_b12x_kda_prefill() + assert api is not None + assert NULL_BLOCK_ID == 0 + num_seqs = cu_seqlens.numel() - 1 + tokens, heads = q.shape[1], q.shape[2] + pool = torch.zeros( + (num_seqs + 1, *initial_state.shape[1:]), + dtype=initial_state.dtype, + device=q.device, + ) + pool[1:] = initial_state + final_indices = torch.arange(1, num_seqs + 1, dtype=torch.int32, device=q.device) + initial_indices = final_indices.masked_fill(~has_initial_state, NULL_BLOCK_ID) + tensors = { + "q": q[0], + "k": k[0], + "v": v[0], + "raw_g": raw_g[0], + "raw_beta": raw_beta[0], + "A_log": A_log, + "dt_bias": dt_bias, + "recurrent_state": pool, + "cu_seqlens": cu_seqlens, + "initial_state_indices": initial_indices, + "final_state_indices": final_indices, + "checkpoint_state_indices": torch.full_like(final_indices, NULL_BLOCK_ID), + "checkpoint_offsets": torch.zeros_like(final_indices), + "num_seqs": torch.tensor([num_seqs], dtype=torch.int32, device=q.device), + "num_tokens": torch.tensor([tokens], dtype=torch.int32, device=q.device), + "output": torch.empty_like(v[0]), + } + plan = api.plan( + api.Caps( + device=q.device, + max_tokens=tokens, + max_seqs=num_seqs, + max_state_slots=num_seqs + 1, + heads=heads, + null_state_index=NULL_BLOCK_ID, + ), + invocation=api.invocation_from_tensors( + A_log=A_log, dt_bias=dt_bias, initial_state_indices=initial_indices + ), + ) + original_pool = pool.clone() + + def prepare_call(state): + (spec,) = state.layout.scratch_specs() + binding = state.bind( + scratch=torch.empty(spec.shape, dtype=spec.dtype, device=q.device), + **tensors, + ) + + def reset(): + pool.copy_(original_pool) + + return PreparedCall( + run=lambda: state.run(binding, lower_bound=lower_bound), + reset=reset, + restore=reset, + ) + + session = PreparationSession(device=q.device, autotune=False) + result = session.prepare( + (plan.request(name="kimi-k3-kda-prefill", prepare_call=prepare_call),) + ) + try: + (spec,) = require_prepared(plan, "sequence.kda_prefill").layout.scratch_specs() + pool.copy_(original_pool) + binding = api.bind( + plan, + scratch=torch.empty(spec.shape, dtype=spec.dtype, device=q.device), + **tensors, + ) + api.run(binding, lower_bound=lower_bound) + torch.accelerator.synchronize() + finally: + result.close() + session.close() + return tensors["output"].unsqueeze(0), pool[1:].clone() + + def _kda_prefill_reference( inputs: SimpleNamespace, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1477,6 +1604,7 @@ def _kda_prefill_reference( [ pytest.param("flashkda", torch.float32, id="flashkda"), pytest.param("flashinfer", torch.bfloat16, id="flashinfer"), + pytest.param("b12x", torch.float32, id="b12x"), ], ) @torch.inference_mode() @@ -1524,6 +1652,7 @@ def test_kda_prefill_near_collinear_keys_remain_finite( pytest.param("flashinfer", torch.bfloat16, 0.03, id="flashinfer-bf16"), pytest.param("flashkda", torch.bfloat16, 0.03, id="flashkda-bf16"), pytest.param("flashkda", torch.float32, 0.01, id="flashkda-fp32"), + pytest.param("b12x", torch.float32, 0.01, id="b12x-fp32"), ], ) @torch.inference_mode() @@ -1545,6 +1674,69 @@ def test_kda_prefill_correctness( assert_close("ht", expected_state, actual_state, tolerance) +@torch.inference_mode() +def test_b12x_kda_prefill_matches_triton_serving_batch(): + """Compare b12x with Triton on a TP9-shaped 4,096-token prefill batch. + + Eleven local heads, ragged requests including a one-token row, and one + request without a computed prefix; b12x reads and writes state slots in + place while Triton consumes gathered initial states. + """ + lower_bound = -5.0 + _require_kda_prefill_backend("b12x", torch.float32, lower_bound) + torch.manual_seed(7) + H, D = 11, 128 + lengths = [1, 700, 1531, 1864] + T = sum(lengths) + q, k, v, raw_g = [ + torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4) + ] + raw_beta = torch.randn(1, T, H, dtype=torch.bfloat16, device=DEVICE) + A_log = torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5 + dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE) * 0.1 + initial_state = ( + torch.randn(len(lengths), H, D, D, dtype=torch.float32, device=DEVICE) * 0.1 + ) + has_initial_state = torch.tensor( + [True, False, True, True], dtype=torch.bool, device=DEVICE + ) + cu_seqlens = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEVICE) + cu_seqlens[1:] = torch.tensor(lengths, device=DEVICE).cumsum(0) + + b12x_out, b12x_state = _b12x_kda_prefill( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + has_initial_state=has_initial_state, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + triton_out, triton_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + g_bias=dt_bias.reshape(-1), + lower_bound=lower_bound, + initial_state=initial_state * has_initial_state[:, None, None, None], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ) + + assert torch.isfinite(b12x_out).all() + assert torch.isfinite(b12x_state).all() + assert_close("o", triton_out, b12x_out, 0.01) + assert_close("ht", triton_state, b12x_state, 0.01) + + @torch.inference_mode() def test_flashinfer_kda_prefill_breakable_graph_cross_stream(): if not is_flashinfer_recurrent_kda_prefill_supported( diff --git a/tests/v1/attention/test_b12x.py b/tests/v1/attention/test_b12x.py index cc7e105c8429..1a581903537a 100644 --- a/tests/v1/attention/test_b12x.py +++ b/tests/v1/attention/test_b12x.py @@ -41,12 +41,14 @@ ) @pytest.mark.parametrize("parallel_drafting", [False, True]) @pytest.mark.parametrize("page_padding", [0, 192]) +@pytest.mark.parametrize("partial_dtype", ["bf16", "fp32"]) @torch.inference_mode() def test_b12x_dense_mla_prepared_capacity_replay_and_high_pages( - dtype, dcp_size, parallel_drafting, page_padding, monkeypatch + dtype, dcp_size, parallel_drafting, page_padding, partial_dtype, monkeypatch ): """Real prepared kernels consume high page IDs and mutable graph inputs.""" _require_b12x_paged_attention() + monkeypatch.setenv("VLLM_K3_DENSE_MLA_PARTIAL_DTYPE", partial_dtype) from b12x._lib.runtime_control import kernel_resolution_guard from b12x.attention import dense_mla from b12x.preparation import PreparationSession diff --git a/vllm/envs.py b/vllm/envs.py index 5a50d57da8fb..76fbe397aa49 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -172,6 +172,7 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_MLA_DISABLE: bool = False + VLLM_K3_DENSE_MLA_PARTIAL_DTYPE: Literal["bf16", "fp32"] = "bf16" VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" VLLM_CUDART_SO_PATH: str | None = None @@ -1499,6 +1500,10 @@ def _resolve_rust_cli_path() -> str | None: ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), + # Keep dense MLA split partials in FP32 until the merge when requested. + "VLLM_K3_DENSE_MLA_PARTIAL_DTYPE": env_with_choices( + "VLLM_K3_DENSE_MLA_PARTIAL_DTYPE", "bf16", ["bf16", "fp32"] + ), # If set, vLLM will pick up the provided Flash Attention MLA # Number of GPUs per worker in Ray, if it is set to be a fraction, # it allows ray to schedule multiple actors on a single GPU, diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index af7a25806b03..e43cbf77dfc0 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import replace +from typing import Any import torch from einops import rearrange @@ -60,6 +61,13 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig from vllm.triton_utils import tl, triton +from vllm.utils.b12x import ( + B12xPreparationUnit, + B12xWorkload, + PreparationResourceUnavailableError, + get_b12x_kda_prefill, + set_b12x_preparation_provider, +) from vllm.utils.flashinfer import ( flashinfer_fused_kda_decode, flashinfer_recurrent_kda, @@ -328,6 +336,24 @@ def is_flashkda_supported( ) +def is_b12x_kda_prefill_supported( + head_dim: int, + input_dtype: torch.dtype, + recurrent_state_dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + api = get_b12x_kda_prefill() + if api is None or not current_platform.is_cuda(): + return False + return ( + head_dim == 128 + and input_dtype == torch.bfloat16 + and recurrent_state_dtype == torch.float32 + and lower_bound is not None + and api.is_supported(torch.device(current_platform.current_device())) + ) + + def is_flashinfer_recurrent_kda_prefill_supported( head_dim: int, input_dtype: torch.dtype, @@ -501,8 +527,26 @@ def resolve_kda_prefill_backend( recurrent_state_dtype: torch.dtype, lower_bound: float | None, ) -> str: - if backend not in ("auto", "triton", "flashkda", "flashinfer"): + """Resolve the packed KDA prefill implementation for one server. + + ``auto`` never selects ``b12x``; that backend must be requested by name. + """ + if backend not in ("auto", "triton", "flashkda", "flashinfer", "b12x"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") + if backend == "b12x": + if not is_b12x_kda_prefill_supported( + head_dim, + input_dtype, + recurrent_state_dtype, + lower_bound, + ): + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package on a " + "supported CUDA device, bfloat16 activations, head_dim=128, " + "float32 recurrent state, and a bounded KDA gate." + ) + logger.info_once("Using b12x KDA prefill backend.") + return "b12x" flashinfer_supported = is_flashinfer_recurrent_kda_prefill_supported( head_dim, input_dtype, @@ -872,6 +916,9 @@ def __init__( (1, T, H, D), self.model_config.dtype, ) + self._b12x_prefill_api: Any | None = None + self._b12x_prefill_plan = None + self._initialize_b12x_kda_prefill(vllm_config) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None @@ -911,11 +958,327 @@ def __init__( self.gemm_rs_ar = gemm_rs_ar else: gemm_rs_ar.warn_incompatible_projection() + self._b12x_preparation_prefix = prefix + if self._b12x_prefill_api is not None: + set_b12x_preparation_provider(self, self) compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + super().bind_kv_cache(kv_cache) + self._b12x_prefill_plan = None + if self._b12x_prefill_api is not None: + self._b12x_prefill_plan = self._b12x_kda_prefill_declaration( + self.kv_cache[1].shape[0] + ) + + def unbind_kv_cache(self) -> None: + # Only bindings to this state-pool generation become invalid. + self._b12x_prefill_plan = None + super().unbind_kv_cache() + + def _initialize_b12x_kda_prefill(self, vllm_config: VllmConfig) -> None: + if self.kda_prefill_backend != "b12x": + return + api = get_b12x_kda_prefill() + if api is None: + raise RuntimeError( + "The b12x KDA prefill backend requires the b12x package." + ) + device = torch.device(current_platform.current_device()) + scheduler_config = vllm_config.scheduler_config + self._b12x_prefill_api = api + self._b12x_prefill_max_tokens = int(scheduler_config.max_num_batched_tokens) + self._b12x_prefill_max_seqs = int(scheduler_config.max_num_seqs) + max_seqs = self._b12x_prefill_max_seqs + # Device-side metadata the op reads by address; one set per layer. + for name, fill in ( + ("_b12x_prefill_cu_seqlens", None), + ("_b12x_prefill_initial_indices", None), + ("_b12x_prefill_null_indices", NULL_BLOCK_ID), + ("_b12x_prefill_zero_offsets", None), + ): + size = max_seqs + 1 if name == "_b12x_prefill_cu_seqlens" else max_seqs + buffer = torch.zeros(size, dtype=torch.int32, device=device) + if fill is not None: + buffer.fill_(fill) + self.register_buffer(name, buffer, persistent=False) + for name in ("_b12x_prefill_num_seqs", "_b12x_prefill_num_tokens"): + self.register_buffer( + name, + torch.zeros(1, dtype=torch.int32, device=device), + persistent=False, + ) + + def _b12x_kda_prefill_declaration(self, max_state_slots: int): + api = self._b12x_prefill_api + if api is None: + raise RuntimeError("b12x KDA prefill was not initialized") + return api.plan( + api.Caps( + device=current_platform.current_device(), + max_tokens=self._b12x_prefill_max_tokens, + max_seqs=self._b12x_prefill_max_seqs, + max_state_slots=max_state_slots, + heads=self.local_num_heads, + head_dim=self.head_dim, + model_dtype=self.model_config.dtype, + state_dtype=self.get_state_dtype()[1], + qk_l2norm=True, + # Prefill chunks end on recurrent block boundaries, as with + # Triton; no mid-chunk checkpoint slot is written. + checkpoint_export=False, + null_state_index=NULL_BLOCK_ID, + ), + invocation=api.invocation_from_tensors( + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state_indices=self._b12x_prefill_initial_indices, + ), + ) + + def _b12x_prefill_scratch_spec(self) -> tuple[tuple[int, ...], torch.dtype]: + plan = self._b12x_prefill_plan + if plan is None: + raise PreparationResourceUnavailableError( + "b12x KDA prefill is not prepared for the current KV generation" + ) + (spec,) = plan.scratch_specs() + return spec.shape, spec.dtype + + def get_b12x_preparation_units( + self, + layer: torch.nn.Module, + workload: B12xWorkload, + ) -> tuple[B12xPreparationUnit, ...]: + if layer is not self: + raise ValueError("KDA preparation owner mismatch") + if self._b12x_prefill_plan is None: + return () + request = self._b12x_prefill_plan.request( + name=f"{self._b12x_preparation_prefix}.kda.prefill", + prepare_call=self._prepare_b12x_kda_prefill, + benchmark_call=self._benchmark_b12x_kda_prefill, + ) + return ( + B12xPreparationUnit( + name="KDA prefill", + key=(self._b12x_preparation_prefix, "kda-prefill"), + requests=(request,), + stage="state", + ), + ) + + @staticmethod + def _benchmark_values(tensor: torch.Tensor) -> None: + values = torch.arange( + tensor.numel(), dtype=tensor.dtype, device=tensor.device + ).reshape_as(tensor) + tensor.copy_(values.div_(max(values.numel(), 1))) + + @staticmethod + def _b12x_trial_tensor(source: torch.Tensor) -> torch.Tensor: + """Allocate a trial tensor with the source's layout and alignment.""" + storage_size = 1 + sum( + (size - 1) * stride + for size, stride in zip(source.shape, source.stride()) + if size + ) + source_pointer = source.data_ptr() + source_alignment = ( + min(16, source_pointer & -source_pointer) if source_pointer else 16 + ) + storage = torch.empty( + storage_size + 16, dtype=source.dtype, device=source.device + ) + for storage_offset in range(16): + trial = storage.as_strided( + source.shape, source.stride(), storage_offset=storage_offset + ) + pointer = trial.data_ptr() + alignment = min(16, pointer & -pointer) if pointer else 16 + if alignment == source_alignment: + return trial + return storage.as_strided(source.shape, source.stride()) + + def _prepare_b12x_kda_prefill(self, state): + return self._b12x_kda_prefill_call(state, benchmark=False) + + def _benchmark_b12x_kda_prefill(self, state): + return self._b12x_kda_prefill_call(state, benchmark=True) + + def _b12x_kda_prefill_call(self, state, *, benchmark: bool): + from b12x.preparation import PreparedCall + + specs = tuple(state.layout.scratch_specs()) + if len(specs) != 1: + raise RuntimeError("b12x KDA prefill requires exactly one scratch buffer") + slots = self.kv_cache[1] + slot = 1 if slots.shape[0] > 1 else 0 + # Preparation owns its scratch and synthetic activations; serving + # binds the live request tensors and workspace-manager scratch. + scratch = torch.empty(specs[0].shape, dtype=specs[0].dtype, device=slots.device) + shape = (self._b12x_prefill_max_tokens, self.local_num_heads, self.head_dim) + q, k, v, raw_g, output = ( + torch.empty(shape, dtype=self.model_config.dtype, device=slots.device) + for _ in range(5) + ) + raw_beta = torch.empty( + shape[:2], dtype=self.model_config.dtype, device=slots.device + ) + if benchmark: + cu_seqlens = self._b12x_trial_tensor(self._b12x_prefill_cu_seqlens) + indices = self._b12x_trial_tensor(self._b12x_prefill_initial_indices) + checkpoint_indices = self._b12x_trial_tensor( + self._b12x_prefill_null_indices + ) + offsets = self._b12x_trial_tensor(self._b12x_prefill_zero_offsets) + num_seqs = self._b12x_trial_tensor(self._b12x_prefill_num_seqs) + num_tokens = self._b12x_trial_tensor(self._b12x_prefill_num_tokens) + saved_state = slots[slot : slot + 1].clone() + else: + cu_seqlens = self._b12x_prefill_cu_seqlens + indices = self._b12x_prefill_initial_indices + checkpoint_indices = self._b12x_prefill_null_indices + offsets = self._b12x_prefill_zero_offsets + num_seqs = self._b12x_prefill_num_seqs + num_tokens = self._b12x_prefill_num_tokens + saved_state = None + + def produce(): + for tensor in (q, k, v, raw_g, raw_beta): + self._benchmark_values(tensor) + cu_seqlens.zero_() + cu_seqlens[1:].fill_(self._b12x_prefill_max_tokens) + indices.fill_(slot) + if benchmark: + checkpoint_indices.fill_(NULL_BLOCK_ID) + offsets.zero_() + num_seqs.fill_(1) + num_tokens.fill_(self._b12x_prefill_max_tokens) + + def reset(): + if saved_state is not None: + slots[slot : slot + 1].copy_(saved_state) + + binding = state.bind( + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias.view(-1, self.head_dim), + recurrent_state=slots, + cu_seqlens=cu_seqlens, + initial_state_indices=indices, + final_state_indices=indices, + checkpoint_state_indices=checkpoint_indices, + checkpoint_offsets=offsets, + num_seqs=num_seqs, + num_tokens=num_tokens, + output=output, + ) + return PreparedCall( + run=lambda: state.run( + binding, + lower_bound=self.gate_lower_bound, + max_live_tokens=self._b12x_prefill_max_tokens, + max_live_seqs=1, + ), + produce=produce, + reset=reset, + restore=reset if saved_state is not None else None, + ) + + def _run_b12x_kda_prefill( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + cu_seqlens: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + recurrent_state: torch.Tensor, + output: torch.Tensor | None, + ) -> torch.Tensor: + """Advance the recurrent-state pool in place and return the output. + + Without a caller ``output``, the rows land in workspace storage + reserved together with the scratch. + """ + api = self._b12x_prefill_api + plan = self._b12x_prefill_plan + if api is None or plan is None: + raise PreparationResourceUnavailableError( + "b12x KDA prefill is not prepared for the current KV generation" + ) + num_tokens, num_requests = int(q.shape[0]), int(state_indices.shape[0]) + if ( + num_tokens > self._b12x_prefill_max_tokens + or num_requests > self._b12x_prefill_max_seqs + ): + raise ValueError( + "b12x KDA prefill capacity exceeded: " + f"tokens={num_tokens}/{self._b12x_prefill_max_tokens}, " + f"requests={num_requests}/{self._b12x_prefill_max_seqs}" + ) + manager = current_workspace_manager() + if output is None: + scratch, output = manager.get_simultaneous( + self._b12x_prefill_scratch_spec(), + ( + ( + self._b12x_prefill_max_tokens, + self.local_num_heads, + self.head_dim, + ), + self.model_config.dtype, + ), + ) + output = output[:num_tokens] + else: + (scratch,) = manager.get_simultaneous(self._b12x_prefill_scratch_spec()) + # A request without a computed prefix starts from the zero state. + initial_indices = self._b12x_prefill_initial_indices[:num_requests] + initial_indices.copy_(state_indices) + initial_indices.masked_fill_(~has_initial_state[:num_requests], NULL_BLOCK_ID) + self._b12x_prefill_num_seqs.fill_(num_requests) + self._b12x_prefill_num_tokens.fill_(num_tokens) + binding = api.bind( + plan, + scratch=scratch, + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=self.A_log, + dt_bias=self.dt_bias.view(-1, self.head_dim), + recurrent_state=recurrent_state, + cu_seqlens=cu_seqlens[: num_requests + 1], + initial_state_indices=initial_indices, + final_state_indices=state_indices, + checkpoint_state_indices=self._b12x_prefill_null_indices[:num_requests], + checkpoint_offsets=self._b12x_prefill_zero_offsets[:num_requests], + num_seqs=self._b12x_prefill_num_seqs, + num_tokens=self._b12x_prefill_num_tokens, + output=output, + ) + api.run( + binding, + lower_bound=self.gate_lower_bound, + max_live_tokens=num_tokens, + max_live_seqs=num_requests, + ) + return output + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: spec = super().get_kv_cache_spec(vllm_config) assert isinstance(spec, MambaSpec) @@ -1297,12 +1660,38 @@ def _prefill_conv( assert non_spec_state_indices_tensor is not None assert has_initial_state is not None - initial_state = gather_initial_states( - recurrent_state, - non_spec_state_indices_tensor, - has_initial_state, + # The b12x op reads and writes the state pool in place. + initial_state = ( + None + if self.kda_prefill_backend == "b12x" + else gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) ) - if self.kda_prefill_backend == "flashkda": + if self.kda_prefill_backend == "b12x": + assert self.gate_lower_bound is not None + assert non_spec_query_start_loc is not None + b12x_out = None + if non_spec_out is not None: + b12x_out = non_spec_out[0] + elif not has_spec_decode: + b12x_out = core_attn_out[0, : q_ns.shape[1]] + core_attn_out_non_spec = self._run_b12x_kda_prefill( + q=q_ns[0], + k=k_ns[0], + v=v_ns[0], + raw_g=g1_ns[0], + raw_beta=beta_ns[0], + cu_seqlens=non_spec_query_start_loc, + state_indices=non_spec_state_indices_tensor, + has_initial_state=has_initial_state, + recurrent_state=recurrent_state, + output=b12x_out, + ).unsqueeze(0) + elif self.kda_prefill_backend == "flashkda": + assert initial_state is not None assert self.gate_lower_bound is not None assert self._flashkda_buffer_specs is not None workspace_out, final_state, checkpoint_state, workspace = ( @@ -1395,6 +1784,7 @@ def _prefill_conv( workspace=workspace, ) elif self.kda_prefill_backend == "flashinfer": + assert initial_state is not None assert self.gate_lower_bound is not None assert m.flashinfer_prefill_query_start_loc is not None if q_ns.shape[1] > initial_state.shape[0]: @@ -1444,9 +1834,10 @@ def _prefill_conv( cu_seqlens=non_spec_query_start_loc, out=non_spec_out, ) - recurrent_state[non_spec_state_indices_tensor] = ( - last_recurrent_state.to(recurrent_state.dtype) - ) + if self.kda_prefill_backend != "b12x": + recurrent_state[non_spec_state_indices_tensor] = ( + last_recurrent_state.to(recurrent_state.dtype) + ) else: # Pure non-speculative decode. assert non_spec_state_indices_tensor is not None @@ -1490,7 +1881,7 @@ def _prefill_conv( ) elif core_attn_out_non_spec is not None: if ( - self.kda_prefill_backend not in ("flashkda", "flashinfer") + self.kda_prefill_backend not in ("flashkda", "flashinfer", "b12x") or m.num_prefills == 0 ): # TODO: decode kernels write directly to core_attn_out diff --git a/vllm/v1/attention/backends/mla/b12x_mla.py b/vllm/v1/attention/backends/mla/b12x_mla.py index e9929439fdfd..0b8bb159bb14 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla.py +++ b/vllm/v1/attention/backends/mla/b12x_mla.py @@ -7,6 +7,7 @@ import torch +from vllm import envs from vllm.config import get_current_vllm_config from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group @@ -342,6 +343,11 @@ def get_b12x_preparation_units(self, layer, workload: B12xWorkload): v_head_dim=self.kv_lora_rank, physical_record_width=cache.shape[2], use_cuda_graph=True, + partial_dtype=( + torch.float32 + if envs.VLLM_K3_DENSE_MLA_PARTIAL_DTYPE == "fp32" + else torch.bfloat16 + ), ) def descriptor(shape, strides, dtype):