Add DCP + Eagle support for Tokenspeed MLA backends (#48180)
Signed-off-by: Pavani Majety <pmajety@nvidia.com> Signed-off-by: Jingyi Yang <girasoleyang@gmail.com> Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
This commit is contained in:
Pavani Majety
co-authored by
Jingyi Yang
Jingyi Yang
parent
18c4067a54
commit
7fc97042c3
@@ -232,7 +232,7 @@ MLA decode backends are selected using the standard
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2
|
||||
quack-kernels>=0.3.3
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.2; platform_system == "Linux"
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
# Humming kernels for quantization gemm
|
||||
humming-kernels[cu13]==0.1.10
|
||||
|
||||
@@ -154,6 +154,17 @@ def test_use_dcp_fallback(_mock):
|
||||
assert _call(dcp_world_size=2) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_dcp_fallback_prefill(_mock):
|
||||
# TRTLLM prefill has no cross-rank LSE combine for DCP-sharded KV.
|
||||
assert _call(dcp_world_size=2, is_prefill=True) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
|
||||
def test_use_dcp_fallback_prefill_force_on_still_false(_mock):
|
||||
assert _call(dcp_world_size=2, is_prefill=True, force_use_trtllm=True) is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
|
||||
def test_use_platform_unsupported(_mock):
|
||||
assert _call() is False
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""FlashInfer GQA builder: reorder threshold under DCP with spec decode."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.utils import create_vllm_config
|
||||
from vllm.config import SpeculativeConfig, set_current_vllm_config
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends import flashinfer as flashinfer_backend
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
FlashInferDecodeKernel,
|
||||
FlashInferMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import PerLayerParameters
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("FlashInfer backend requires a CUDA platform.", allow_module_level=True)
|
||||
|
||||
|
||||
def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch):
|
||||
"""trtllm-gen decode receives no cp_rank/global-seq-len information, so its
|
||||
end-aligned causal mask is wrong for q_len > 1 over the DCP-interleaved
|
||||
local KV shard. The builder must keep reorder_batch_threshold at 1 under
|
||||
DCP so spec queries take the (DCP-aware) prefill path instead.
|
||||
"""
|
||||
vllm_config = create_vllm_config(max_model_len=1024)
|
||||
vllm_config.parallel_config.decode_context_parallel_size = 2
|
||||
vllm_config.speculative_config = SpeculativeConfig(
|
||||
method="ngram", num_speculative_tokens=3
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
flashinfer_backend, "can_use_trtllm_attention", lambda *args, **kwargs: True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
FlashInferMetadataBuilder,
|
||||
"_get_flashinfer_trtllm_api_decode_kernel",
|
||||
staticmethod(lambda: FlashInferDecodeKernel.TRTLLM_GEN),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
flashinfer_backend,
|
||||
"get_per_layer_parameters",
|
||||
lambda *args, **kwargs: {
|
||||
"layer.0": PerLayerParameters(
|
||||
window_left=-1, logits_soft_cap=None, sm_scale=0.1, has_sinks=False
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
kv_cache_spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=vllm_config.model_config.get_num_kv_heads(
|
||||
vllm_config.parallel_config
|
||||
),
|
||||
head_size=vllm_config.model_config.get_head_size(),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
builder = FlashInferMetadataBuilder(
|
||||
kv_cache_spec,
|
||||
["layer.0"],
|
||||
vllm_config,
|
||||
torch.device("cpu"),
|
||||
)
|
||||
|
||||
# Guard against passing vacuously with the kernel disabled.
|
||||
assert (
|
||||
builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN
|
||||
)
|
||||
assert builder.reorder_batch_threshold == 1
|
||||
@@ -8,6 +8,9 @@ Known Issues:
|
||||
test_backend_correctness[small_prefill], but passes when run alone.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -32,6 +35,7 @@ from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.v1.attention.backend import CommonAttentionMetadata
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
from vllm.v1.attention.backends.mla import flashmla as flashmla_module
|
||||
from vllm.v1.attention.backends.mla import tokenspeed_mla as tokenspeed_mla_module
|
||||
from vllm.v1.attention.backends.mla.prefill import (
|
||||
MLAPrefillBackendEnum,
|
||||
get_mla_prefill_backend,
|
||||
@@ -693,6 +697,100 @@ def test_mock_mla_dcp_fp8_decode_gathers_quantized_query(
|
||||
)
|
||||
|
||||
|
||||
def test_tokenspeed_mla_dcp_single_token_decode_contract(monkeypatch):
|
||||
decode_call = None
|
||||
num_decodes = 2
|
||||
tokens_per_decode = 1
|
||||
num_decode_tokens = num_decodes * tokens_per_decode
|
||||
dcp_world_size = 2
|
||||
dcp_rank = 1
|
||||
num_heads = 128
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
head_size = kv_lora_rank + qk_rope_head_dim
|
||||
block_size = 64
|
||||
num_blocks = 4
|
||||
max_seq_len = 24
|
||||
|
||||
def fake_decode(**kwargs):
|
||||
nonlocal decode_call
|
||||
decode_call = kwargs
|
||||
q = kwargs["query"]
|
||||
out = torch.empty(
|
||||
q.shape[0],
|
||||
q.shape[1],
|
||||
q.shape[2],
|
||||
kv_lora_rank,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
lse = torch.empty(q.shape[0], q.shape[1], q.shape[2], dtype=torch.float32)
|
||||
return out, lse
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"tokenspeed_mla",
|
||||
SimpleNamespace(tokenspeed_mla_decode=fake_decode),
|
||||
)
|
||||
|
||||
impl = object.__new__(tokenspeed_mla_module.TokenspeedMLAImpl)
|
||||
impl.dcp_world_size = dcp_world_size
|
||||
impl.dcp_rank = dcp_rank
|
||||
impl.cp_kv_cache_interleave_size = 1
|
||||
impl.need_to_return_lse_for_decode = True
|
||||
impl.kv_lora_rank = kv_lora_rank
|
||||
impl.qk_rope_head_dim = qk_rope_head_dim
|
||||
impl.num_heads = num_heads
|
||||
impl.scale = 1.0
|
||||
impl.softmax_scale = 1.0
|
||||
impl.output_scale = 1.0
|
||||
impl._workspace_buffer = torch.empty(1, dtype=torch.int8)
|
||||
|
||||
metadata = SimpleNamespace(
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
max_seq_len=max_seq_len,
|
||||
decode=SimpleNamespace(
|
||||
block_table=torch.empty((num_decodes, 1), dtype=torch.int32),
|
||||
seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
|
||||
dcp_tot_seq_lens=torch.tensor([16, max_seq_len], dtype=torch.int32),
|
||||
),
|
||||
)
|
||||
q = torch.empty(num_decode_tokens, num_heads, head_size, dtype=torch.float8_e4m3fn)
|
||||
kv_cache = torch.empty(
|
||||
num_blocks,
|
||||
block_size,
|
||||
head_size,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
)
|
||||
|
||||
out, lse = impl.forward_mqa(
|
||||
q,
|
||||
kv_cache,
|
||||
metadata,
|
||||
SimpleNamespace(_q_scale_float=2.0, _k_scale_float=3.0),
|
||||
)
|
||||
|
||||
assert out.shape == (num_decode_tokens, num_heads, kv_lora_rank)
|
||||
assert lse is not None
|
||||
assert lse.shape == (num_decode_tokens, num_heads)
|
||||
|
||||
assert decode_call is not None
|
||||
assert decode_call["query"].shape == (
|
||||
num_decodes,
|
||||
tokens_per_decode,
|
||||
num_heads,
|
||||
head_size,
|
||||
)
|
||||
torch.testing.assert_close(decode_call["seq_lens"], metadata.decode.seq_lens)
|
||||
torch.testing.assert_close(decode_call["block_tables"], metadata.decode.block_table)
|
||||
torch.testing.assert_close(
|
||||
decode_call["causal_seqs"], metadata.decode.dcp_tot_seq_lens
|
||||
)
|
||||
assert decode_call["return_lse"] is True
|
||||
assert decode_call["cp_world"] == dcp_world_size
|
||||
assert decode_call["cp_rank"] == dcp_rank
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"])
|
||||
def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads(
|
||||
monkeypatch, is_fp8_kvcache
|
||||
|
||||
@@ -439,11 +439,13 @@ def use_trtllm_attention(
|
||||
if force_use_trtllm is not None and not force_use_trtllm:
|
||||
return False
|
||||
|
||||
# Decode context parallel is not supported
|
||||
# TRTLLM prefill attends only the DCP-local KV shard and has no
|
||||
# cross-rank LSE combine, so it cannot be used with DCP; fall back to
|
||||
# FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected
|
||||
# separately (all-gathered query heads + LSE combine in forward).
|
||||
if dcp_world_size > 1:
|
||||
logger.warning_once(
|
||||
"Trtllm does not support returning LSE and as a result "
|
||||
"does not support DCP, reverting to FlashInfer"
|
||||
"TRTLLM prefill does not support DCP, reverting to FlashInfer"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
@@ -752,12 +752,32 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
if can_use_xqa_or_trtllm_gen_decode
|
||||
else None
|
||||
)
|
||||
if (
|
||||
self.use_dcp
|
||||
and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
|
||||
):
|
||||
logger.warning_once(
|
||||
"FlashInfer XQA decode does not support returning LSE and "
|
||||
"therefore does not support DCP, reverting to native FlashInfer "
|
||||
"decode."
|
||||
)
|
||||
self.use_trtllm_decode_attention = False
|
||||
self.flashinfer_trtllm_api_decode_kernel = None
|
||||
supports_spec_as_decode = (
|
||||
self.flashinfer_trtllm_api_decode_kernel
|
||||
== FlashInferDecodeKernel.TRTLLM_GEN
|
||||
)
|
||||
self._init_reorder_batch_threshold(
|
||||
1, supports_spec_as_decode=supports_spec_as_decode
|
||||
1,
|
||||
supports_spec_as_decode=supports_spec_as_decode,
|
||||
# trtllm-gen decode receives no cp_rank/global-seq-len information,
|
||||
# so its end-aligned causal mask is wrong for q_len > 1 over the
|
||||
# DCP-interleaved local KV shard (spec token i misses up to
|
||||
# (dcp_world_size - 1) * (q_len - 1 - i) KV entries, including its
|
||||
# own). Keep the threshold at 1 under DCP so spec queries take the
|
||||
# DCP-aware prefill path, until the kernel is CP-aware (compare
|
||||
# flash_attn_varlen_func's cp_world_size/cp_rank/cp_tot_seqused_k).
|
||||
supports_dcp_with_varlen=False,
|
||||
)
|
||||
|
||||
self._cascade_wrapper = None # Wrapper for cascade attention
|
||||
@@ -1127,9 +1147,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
has_sinks=self.has_sinks,
|
||||
has_spec=uses_spec_reorder,
|
||||
)
|
||||
decode_with_flashinfer_trtllm_api = (
|
||||
causal and self.use_trtllm_decode_attention and self.dcp_world_size <= 1
|
||||
)
|
||||
decode_with_flashinfer_trtllm_api = causal and self.use_trtllm_decode_attention
|
||||
|
||||
if not causal and self.use_dcp:
|
||||
raise NotImplementedError(
|
||||
@@ -1230,7 +1248,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
# Compute paged_kv_indices if necessary
|
||||
# paged_kv_indices is only needed for FlashInfer native paths;
|
||||
# XQA/trtllm-gen paths use block_tables directly on GPU.
|
||||
needs_paged_kv_indices = use_cascade or not all_uses_trtllm
|
||||
needs_native_paged_prefill = num_prefills > 0 and not prefill_use_trtllm
|
||||
needs_native_paged_decode = (
|
||||
num_decodes > 0 and not decode_with_flashinfer_trtllm_api
|
||||
)
|
||||
needs_paged_kv_indices = (
|
||||
use_cascade or needs_native_paged_prefill or needs_native_paged_decode
|
||||
)
|
||||
if needs_paged_kv_indices:
|
||||
assert num_blocks_np is not None
|
||||
assert seq_lens_np is not None
|
||||
@@ -1306,6 +1330,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
assert qo_indptr_prefill_cpu.shape[0] == num_prefills + 1
|
||||
|
||||
if prefill_use_trtllm:
|
||||
# TRTLLM prefill has no cross-rank combine for DCP-sharded KV;
|
||||
# use_trtllm_attention never selects it when DCP is enabled.
|
||||
assert not self.use_dcp
|
||||
# Create GPU versions
|
||||
qo_indptr_prefill_gpu = (
|
||||
qo_indptr[prefill_start:] - qo_indptr[prefill_start]
|
||||
@@ -1332,7 +1359,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
|
||||
attn_metadata.prefill = TRTLLMPrefill(
|
||||
block_tables=block_table_tensor[prefill_start:],
|
||||
seq_lens=seq_lens[prefill_start:],
|
||||
seq_lens=prefill_seq_lens,
|
||||
cum_seq_lens_q=qo_indptr_prefill_gpu,
|
||||
cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,
|
||||
max_q_len=max_q_len_prefill,
|
||||
@@ -1409,10 +1436,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
|
||||
f"Got {num_decode_tokens=} and {num_decodes=}."
|
||||
)
|
||||
assert self.flashinfer_trtllm_api_decode_kernel is not None
|
||||
seq_lens_decode = seq_lens[:num_decodes]
|
||||
if self.use_dcp:
|
||||
assert common_attn_metadata.dcp_local_seq_lens is not None
|
||||
seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[
|
||||
:num_decodes
|
||||
]
|
||||
attn_metadata.decode = FlashInferTrtllmAPIDecode(
|
||||
kernel=self.flashinfer_trtllm_api_decode_kernel,
|
||||
block_tables=block_table_tensor[:num_decodes],
|
||||
seq_lens=seq_lens[:num_decodes],
|
||||
seq_lens=seq_lens_decode,
|
||||
max_seq_len=max_seq_len,
|
||||
)
|
||||
else:
|
||||
@@ -2086,6 +2119,18 @@ class FlashInferImpl(AttentionImpl):
|
||||
f"contiguous, got strides {kv_strides}"
|
||||
)
|
||||
|
||||
if use_dcp:
|
||||
assert decode_with_trtllm_gen
|
||||
if output.dtype == FP4_DTYPE:
|
||||
raise NotImplementedError(
|
||||
"DCP decode with FlashInfer trtllm-gen does not support "
|
||||
"FP4 attention output yet."
|
||||
)
|
||||
decode_query = get_dcp_group().all_gather(
|
||||
decode_query.contiguous(), dim=-2
|
||||
)
|
||||
decode_query = canonicalize_singleton_dim_strides(decode_query)
|
||||
|
||||
if output.dtype == FP4_DTYPE:
|
||||
assert self.o_sf_scale is not None
|
||||
out = FP4Tensor(
|
||||
@@ -2124,6 +2169,19 @@ class FlashInferImpl(AttentionImpl):
|
||||
else self.bmm1_scale
|
||||
)
|
||||
|
||||
lse = None
|
||||
if use_dcp:
|
||||
out = torch.empty(
|
||||
decode_query.shape,
|
||||
dtype=output.dtype,
|
||||
device=decode_query.device,
|
||||
)
|
||||
lse = torch.empty(
|
||||
(decode_query.size(0), decode_query.size(1)),
|
||||
dtype=torch.float32,
|
||||
device=decode_query.device,
|
||||
)
|
||||
|
||||
trtllm_batch_decode_with_kv_cache(
|
||||
query=decode_query,
|
||||
kv_cache=(
|
||||
@@ -2145,9 +2203,19 @@ class FlashInferImpl(AttentionImpl):
|
||||
kv_cache_sf=(
|
||||
nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
|
||||
),
|
||||
lse=lse,
|
||||
return_lse=self.need_to_return_lse_for_decode,
|
||||
)
|
||||
|
||||
if needs_fp8_out:
|
||||
if use_dcp:
|
||||
assert isinstance(out, torch.Tensor)
|
||||
assert lse is not None
|
||||
output[:num_decode_tokens] = self.dcp_combine(
|
||||
out,
|
||||
lse,
|
||||
get_dcp_group(),
|
||||
)
|
||||
elif needs_fp8_out:
|
||||
output[:num_decode_tokens].copy_(out.to(output.dtype))
|
||||
return output_padded
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""TokenSpeed CuTe DSL MLA decode backend (Blackwell, FP8 KV cache only)."""
|
||||
|
||||
from typing import ClassVar
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
import torch
|
||||
|
||||
@@ -25,6 +25,10 @@ from vllm.v1.attention.backend import (
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import KVCacheLayoutType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Workspace upper bound for tokenspeed_mla_decode (per-device, lazy):
|
||||
@@ -55,6 +59,22 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
|
||||
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
|
||||
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kv_cache_spec: "AttentionSpec",
|
||||
layer_names: list[str],
|
||||
vllm_config: "VllmConfig",
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
kv_cache_spec,
|
||||
layer_names,
|
||||
vllm_config,
|
||||
device,
|
||||
MLACommonMetadata,
|
||||
supports_dcp_with_varlen=True,
|
||||
)
|
||||
|
||||
|
||||
class TokenspeedMLABackend(MLACommonBackend):
|
||||
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
|
||||
@@ -130,6 +150,11 @@ class TokenspeedMLABackend(MLACommonBackend):
|
||||
|
||||
|
||||
class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
|
||||
can_return_lse_for_decode: bool = True
|
||||
# tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges
|
||||
# partial outputs with exp2(lse).
|
||||
lse_base_on_e: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_heads: int,
|
||||
@@ -228,9 +253,15 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
|
||||
f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}."
|
||||
)
|
||||
|
||||
num_decodes = attn_metadata.num_decodes
|
||||
num_decode_tokens = attn_metadata.num_decode_tokens
|
||||
block_tables = attn_metadata.decode.block_table
|
||||
seq_lens = attn_metadata.decode.seq_lens
|
||||
causal_seqs = attn_metadata.decode.dcp_tot_seq_lens
|
||||
|
||||
# tokenspeed_mla_decode expects query shape
|
||||
# (num_decodes, q_len_per_request, num_heads, head_dim).
|
||||
if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
|
||||
if num_decode_tokens % num_decodes != 0:
|
||||
logger.warning_once(
|
||||
"""TokenspeedMLAImpl got a query of uneven length.
|
||||
This usually indicates an issue in batch reordering
|
||||
@@ -238,7 +269,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
|
||||
)
|
||||
q = q.unsqueeze(1)
|
||||
else:
|
||||
q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
|
||||
q = q.view(num_decodes, -1, q.shape[-2], q.shape[-1])
|
||||
|
||||
if self.softmax_scale is None:
|
||||
# FP8 KV cache is mandatory for this backend, so q_scale/k_scale
|
||||
@@ -257,22 +288,29 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
|
||||
|
||||
# vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
|
||||
# tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
|
||||
o = tokenspeed_mla_decode(
|
||||
return_lse = self.need_to_return_lse_for_decode
|
||||
kernel_out = tokenspeed_mla_decode(
|
||||
query=q,
|
||||
kv_cache=kv_c_and_k_pe_cache,
|
||||
workspace_buffer=self._workspace_buffer,
|
||||
kv_lora_rank=self.kv_lora_rank,
|
||||
qk_rope_head_dim=self.qk_rope_head_dim,
|
||||
block_tables=attn_metadata.decode.block_table,
|
||||
seq_lens=attn_metadata.decode.seq_lens,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=attn_metadata.max_seq_len,
|
||||
softmax_scale=self.softmax_scale,
|
||||
output_scale=self.output_scale,
|
||||
enable_pdl=False,
|
||||
return_lse=return_lse,
|
||||
causal_seqs=causal_seqs if self.dcp_world_size > 1 else None,
|
||||
cp_world=self.dcp_world_size,
|
||||
cp_rank=self.dcp_rank,
|
||||
)
|
||||
if return_lse:
|
||||
o, lse = kernel_out
|
||||
lse = lse.view(-1, lse.shape[-1])
|
||||
else:
|
||||
o, lse = kernel_out, None
|
||||
|
||||
# Flatten the output for consistent shape
|
||||
o = o.view(-1, o.shape[-2], o.shape[-1])
|
||||
|
||||
# tokenspeed_mla_decode does not return LSE.
|
||||
return o, None
|
||||
return o, lse
|
||||
|
||||
Reference in New Issue
Block a user