common.py 72.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
5
# MLA Common Components

6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
This file implements common components for MLA implementations.

First we define:

Sq      as Q sequence length
Skv     as KV sequence length

MLA has two possible ways of computing, a data-movement friendly approach and a
compute friendly approach, we generally want to use the compute friendly
approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1)
and the data-movement friendly approach for "decode" (i.e. the ratio
Sq / Skv is "large").

NOTE what we deem small and large is currently determined by if its labelled
prefill or decode by the scheduler, but this is something we should probably
tune.

Main reference: DeepseekV2 paper, and FlashInfer Implementation
(https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).

Deepseek's MLA attention works the following way:
Matthew Bonanni's avatar
Matthew Bonanni committed
27
* Use a single latent vector to represent the per-token entry of the KV cache.
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
* For decode (i.e. the memory friendly approach) the attention "simulates" a
multi-head attention, while the compute is similar to multi-query attention.

Below is example of both paths assuming batchsize = 1

## More Extent Definitions:

C           Context length, `Skv - Sq`
H           hidden size
N           number of attention heads
Lq          latent dimension for Q              1536 in DSV3
Lkv         latent dimension for K/V            512 in DSV3
P           nope dimension, no rope.            128 in DSV3
R           rope dimension, goes through rope.  64 in DSV3
V           V head dim.                         128 in DSV3

## Vector/Matrix Definitions

h_t         hidden states (input to attention)  shape [Sq, H]
q_c         latent/compressed Q                 shape [Sq, Lq]
q_nope      uncompressed Q (no-rope)            shape [Sq, N, P]
q_pe        uncompressed Q (rope)               shape [Sq, N, R]
kv_c        latent/compressed KV                shape [Skv, Lkv]
k_pe        decoupled k position embeddings     shape [Skv, R]
new_kv_c    new kv_c from current iter          shape [Sq, Lkv]
new_k_pe    new k_pe from current iter          shape [Sq, R]
cache_kv_c  cached k_c from previous iters      shape [C, Lkv]
cache_k_pe  cached k_pe from previous iters     shape [C, R]
W_DQ        project h_t to q_c                  shape [H, Lq]
W_UQ        project q_c to q_nope               shape [Lq, N * P]
W_QR        project q_c to q_pe                 shape [Lq, N * R]
W_DKV       project h_t to kv_c                 shape [H, Lkv]
60
61
62
W_UK        project kv_c to k_nope              shape [Lkv, N, P]
W_KR        project h_t to k_pe                 shape [H, R]
W_UV        project kv_c to v                   shape [Lkv, N, V]
63
64
65
66
67
68
69
70
71
72
73
74
W_O         project v to h_t                    shape [N * V, H]


## Compute Friendly Approach (i.e. "_forward_prefill"):

q_c      = h_t @ W_DQ
q_nope   = (q_c @ W_UQ).view(Sq, N, P)
q_pe     = RoPE(q_c @ W_QR).view(Sq, N, R)
new_kv_c = h_t @ W_DKV
new_k_pe = RoPE(h_t @ W_KR)
kv_c     = torch.cat([new_kv_c, cache_kv_c], dim=0)
k_pe     = torch.cat([new_k_pe, cache_k_pe], dim=0)
75
76
k_nope   = (kv_c @ W_UK.view(Lkv, N * P)).view(Skv, N, P)
v        = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)
77
78
79
80
81
82
83
84

// MHA with QK headdim = P + R
//           V headdim = V
//      spda_o shape [Sq, N, V]
spda_o = scaled_dot_product_attention(
    torch.cat([q_nope, q_pe], dim=-1),
    torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
    v
Matthew Bonanni's avatar
Matthew Bonanni committed
85
)
86
87
88
return spda_o @ W_O

NOTE: in the actual code,
89
90
    `kv_b_proj` is [W_UK; W_UV] concatenated per head
    `q_b_proj` is [W_UQ; W_QR] concatenated per head
91
92
93
94
95
96
97
    `out_proj` is W_O


## Data-Movement Friendly Approach (i.e. "_forward_decode"):

Runtime
q_c      = h_t @ W_DQ
98
99
q_nope   = (q_c @ W_UQ).view(-1, N, P)
ql_nope  = einsum("snh,lnh->snl", q, W_UK)
100
101
102
103
104
105
106
107
108
109
110
111
q_pe     = RoPE(q_c @ W_QR).view(Sq, N, R)
new_kv_c = h_t @ W_DKV
new_k_pe = RoPE(h_t @ W_KR)
kv_c     = torch.cat([new_kv_c, cache_kv_c], dim=0)
k_pe     = torch.cat([new_k_pe, cache_k_pe], dim=0)

// MQA with QK headdim = Lkv + R
//           V headdim = Lkv
//      spda_o shape [Sq, N, Lkv]
// NOTE: this is less compute-friendly since Lkv > P
//       but is more data-movement friendly since its MQA vs MHA
spda_o = scaled_dot_product_attention(
112
    torch.cat([ql_nope, q_pe], dim=-1),
113
114
115
    torch.cat([kv_c, k_pe], dim=-1),
    kv_c
)
116
117
118

o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
return o.view(-1, N * V) @ self.num_heads @ W_O
119
120
121
122


## Chunked Prefill

Matthew Bonanni's avatar
Matthew Bonanni committed
123
124
For chunked prefill we want to use the compute friendly algorithm. We are
assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
125
126
127
128
129
the data-movement friendly approach if the chunk (i.e. `Sq`) is small.

However, the compute-friendly approach can potentially run out of memory if Skv
is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`

Matthew Bonanni's avatar
Matthew Bonanni committed
130
131
To mitigate this, we chunk the computation of attention with respect to the
current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
132
133
134
135
fixed workspace size.

The chunked prefill approach is as follows:

Matthew Bonanni's avatar
Matthew Bonanni committed
136
MCC        Max chunk of context to process per iter, computed dynamically,
137
138
139
140
141
142
143
           used to bound the memory usage

q_c        = h_t @ W_DQ
q_nope     = (q_c @ W_UQ).view(Sq, N, P)
q_pe       = RoPE(q_c @ W_QR).view(Sq, N, R)
new_kv_c   = h_t @ W_DKV
new_k_pe   = RoPE(h_t @ W_KR)
144
145
new_k_nope = (new_kv_c @ W_UK.view(Lkv, N * P)).view(Sq, N, P)
new_v      = (new_kv_c @ W_UV.view(Lkv, N * V)).view(Sq, N, V)
146
147
148
149
150
151
152
153
154
155
156
157

// MHA between queries and new KV
//     with QK headdim = P + R
//           V headdim = V
//    curr_o   shape [Sq, N, V]
//    curr_lse shape [N, Sq], this is just order FA returns
curr_o, curr_lse = scaled_dot_product_attention(
    torch.cat([q_nope, q_pe], dim=-1),
    torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
    new_v,
    casual=True,
    return_softmax_lse=True
Matthew Bonanni's avatar
Matthew Bonanni committed
158
)
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191

// Compute attention with the already existing context
for chunk_idx in range(cdiv(C, MCC)):
    chunk_start  = chunk_idx * MCC
    chunk_end    = min(chunk_start + MCC, C)
    Sc           = chunk_end - chunk_start
    cache_kv_c_chunk   = cache_kv_c[chunk_start:chunk_end]
    cache_k_pe_chunk   = cache_k_pe[chunk_start:chunk_end]
    cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)
    cache_v_chunk      = (cache_kv_c_chunk @ W_UV).view(-1, N, V)

    chunk_o, chunk_lse = scaled_dot_product_attention(
        torch.cat([q_nope, q_pe], dim=-1),
        torch.cat([cache_k_nope_chunk,
                   cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
                   dim=-1),
        cache_v_chunk,
        casual=False,
        return_softmax_lse=True
    )

    curr_o, curr_lse = merge_attn_states(
        suffix_output=curr_o,
        suffix_lse=curr_lse,
        prefix_output=chunk_o,
        prefix_lse=chunk_lse,
    )

return curr_o @ W_O
"""

import functools
from abc import abstractmethod
192
from dataclasses import dataclass, field
193
from typing import ClassVar, Generic, TypeVar
194
195

import torch
196
from tqdm import tqdm
197

198
import vllm.envs as envs
199
from vllm import _custom_ops as ops
200
201
202
203
204
205
from vllm.attention.backends.abstract import (
    AttentionBackend,
    AttentionLayer,
    AttentionMetadata,
    MLAAttentionImpl,
)
206
from vllm.attention.backends.utils import get_mla_dims
207
from vllm.attention.ops.common import cp_lse_ag_out_rs
208
from vllm.attention.ops.merge_attn_states import merge_attn_states
209
from vllm.attention.utils.fa_utils import get_flash_attn_version
210
from vllm.config import VllmConfig, get_current_vllm_config
211
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
212
from vllm.logger import init_logger
213
214
215
216
217
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    LinearBase,
    UnquantizedLinearMethod,
)
Simon Mo's avatar
Simon Mo committed
218
from vllm.platforms import current_platform
219
from vllm.utils import cdiv, round_down
220
from vllm.utils.flashinfer import has_nvidia_artifactory
221
222
223
224
225
226
227
from vllm.v1.attention.backends.utils import (
    AttentionMetadataBuilder,
    CommonAttentionMetadata,
    get_per_layer_parameters,
    infer_global_hyperparameters,
    split_decodes_and_prefills,
)
228
from vllm.v1.kv_cache_interface import AttentionSpec
229
230
231

try:
    from vllm.vllm_flash_attn import flash_attn_varlen_func
232

233
    is_vllm_fa = True
234
235
except ImportError:
    # For rocm use upstream flash attention
236
237
    if current_platform.is_rocm():
        from flash_attn import flash_attn_varlen_func
238
    is_vllm_fa = False
239

240
241
try:
    from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
242
243
    from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache  # noqa: F401

244
245
    flashinfer_available = True
except ImportError:
246
247
    BatchPrefillWithRaggedKVCacheWrapper = object

248
249
    flashinfer_available = False

250
251

def is_rocm_aiter_fp8bmm_enabled() -> bool:
252
253
254
    return (
        current_platform.is_rocm()
        and envs.VLLM_ROCM_USE_AITER_FP8BMM
255
        and envs.VLLM_ROCM_USE_AITER
256
    )
257
258
259


if is_rocm_aiter_fp8bmm_enabled():
260
261
    from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (  # noqa: E501
        batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant as aiter_triton_fp8_bmm,  # noqa: E501
262
    )
263
264

    def dynamic_per_batched_tensor_quant(
265
266
        x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn
    ):
267
268
269
270
271
272
273
274
        DTYPE_MAX = torch.finfo(dtype).max
        min_val, max_val = x.aminmax()
        amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10)
        scale = DTYPE_MAX / amax
        x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX)
        return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()


275
276
logger = init_logger(__name__)

277
278
CUDNN_WORKSPACE_SIZE = 12800

279
280
281
282
283
284

class MLACommonBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
    def get_name() -> str:
285
        return "TRITON_MLA"
286
287

    @staticmethod
288
    def get_metadata_cls() -> type["AttentionMetadata"]:
289
290
291
        return MLACommonMetadata

    @staticmethod
292
    def get_builder_cls() -> type["MLACommonMetadataBuilder"]:
293
294
295
296
297
298
299
300
        return MLACommonMetadataBuilder

    @staticmethod
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,  # assumed to be 1 for MLA
        head_size: int,
301
        cache_dtype_str: str = "auto",
302
    ) -> tuple[int, ...]:
303
304
        return (num_blocks, block_size, head_size)

305
306
307
308
    @classmethod
    def get_supported_dtypes(cls) -> list[torch.dtype]:
        return [torch.float16, torch.bfloat16]

309
310
    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
311
312
        return [576]

313
314
315
316
317
318
319
320
321
    @classmethod
    def validate_head_size(cls, head_size: int) -> None:
        supported_head_sizes = cls.get_supported_head_sizes()
        if head_size not in supported_head_sizes:
            attn_type = cls.__name__.removesuffix("Backend")
            raise ValueError(
                f"Head size {head_size} is not supported by {attn_type}. "
                f"Supported head sizes are: {supported_head_sizes}. "
                "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
322
323
                "FlexAttention backend which supports all head sizes."
            )
324

325
326

@dataclass
327
class MLACommonPrefillMetadata:
328
    """Prefill Specific Metadata"""
329
330
331
332
333
334
335
336
337

    @dataclass
    class ChunkedContextMetadata:
        # New for MLA (compared to FlashAttention)
        # For handling chunked prefill
        cu_seq_lens: torch.Tensor
        starts: torch.Tensor
        seq_tot: list[int]
        max_seq_lens: list[int]
338
        seq_lens: torch.Tensor
339
        workspace: torch.Tensor
340

341
        # for mla DCP
342
343
344
345
346
        cp_chunk_seq_lens: list[list[int]] | None = None
        origin_context_lens: list[int] | None = None
        cp_cu_seq_lens: torch.Tensor | None = None
        chunk_size: int | None = None
        cu_seq_lens_lst: list[list[int]] | None = None
347

348
349
350
    block_table: torch.Tensor
    query_start_loc: torch.Tensor
    max_query_len: int
351
    chunked_context: ChunkedContextMetadata | None = None
352

353

354
355
@dataclass
class FlashInferPrefillMetadata(MLACommonPrefillMetadata):
356
357
    prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
    prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = field(
358
359
        default_factory=list
    )
360
361


362
363
@dataclass
class CudnnPrefillMetadata(MLACommonPrefillMetadata):
364
    class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata):
365
366
        seq_lens: torch.Tensor

367
368
    query_seq_lens: torch.Tensor | None = None
    cudnn_workspace: torch.Tensor | None = None
369
370


371
372
373
374
@dataclass
class MLACommonDecodeMetadata:
    block_table: torch.Tensor
    seq_lens: torch.Tensor
375
    dcp_tot_seq_lens: torch.Tensor | None
376
377
378
379
380
381
382
383
384
385
386
387


D = TypeVar("D", bound=MLACommonDecodeMetadata)


@dataclass
class MLACommonMetadata(Generic[D]):
    """Metadata for MLACommon.

    NOTE: Please read the comment at the top of the file before trying to
    understand this class
    """
388

389
390
391
392
393
394
395
396
    # NOTE(sang): Definition of context_len, query_len, and seq_len.
    # |---------- N-1 iteration --------|
    # |---------------- N iteration ---------------------|
    # |- tokenA -|......................|-- newTokens ---|
    # |---------- context_len ----------|
    # |-------------------- seq_len ---------------------|
    #                                   |-- query_len ---|

397
398
    num_reqs: int
    max_query_len: int
399
    max_seq_len: int
400

401
402
403
404
    num_actual_tokens: int  # Number of tokens excluding padding.
    query_start_loc: torch.Tensor
    slot_mapping: torch.Tensor

405
406
407
408
409
410
    # New for MLA (compared to FlashAttention)
    # For handling prefill decode split
    num_decodes: int
    num_decode_tokens: int
    num_prefills: int

411
    # The dimension of the attention heads
412
    head_dim: int | None = None
413

414
415
416
417
418
419
420
    decode: D | None = None
    prefill: (
        MLACommonPrefillMetadata
        | FlashInferPrefillMetadata
        | CudnnPrefillMetadata
        | None
    ) = None
421
422

    def __post_init__(self):
423
424
        if self.head_dim is not None:
            MLACommonBackend.validate_head_size(self.head_dim)
425
426


427
M = TypeVar("M", bound=MLACommonMetadata)
428
A = TypeVar("A")
429
430


431
def use_flashinfer_prefill() -> bool:
432
    # For blackwell default to flashinfer prefill if it's available since
433
    # it is faster than FA2.
434
435
436
437
438
439
    return (
        not envs.VLLM_DISABLE_FLASHINFER_PREFILL
        and flashinfer_available
        and not envs.VLLM_USE_CUDNN_PREFILL
        and current_platform.is_device_capability(100)
    )
440
441


442
def use_cudnn_prefill() -> bool:
443
444
445
446
447
448
    return (
        flashinfer_available
        and envs.VLLM_USE_CUDNN_PREFILL
        and current_platform.is_device_capability(100)
        and has_nvidia_artifactory()
    )
449
450


451
# Currently 394MB, this can be tuned based on GEMM sizes used.
452
# Chosen to be the same as sglang:
453
454
455
456
#  https://github.com/sgl-project/sglang/blob/766392c6bda2558b61ce6d1c1bfd8081a549e1f1/python/sglang/global_config.py#L37
FLASHINFER_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024


457
class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
458
459
460
461
    """
    NOTE: Please read the comment at the top of the file before trying to
    understand this class
    """
462

463
464
465
466
467
468
469
470
471
472
473
474
475
476
    # Whether the backend supports reordering the batch such that
    # short sequences (i.e. verification for speculative decoding) are
    # classified as decode requests.
    # If True, this will increase `reorder_batch_threshold` (below) when
    # speculative decoding is enabled, and set `require_uniform=True` when
    # when reordering the batch. Non-uniform decode requests will
    # fall back to prefill in this case.
    supports_uniform_spec_as_decode: ClassVar[bool] = False

    # The threshold for reordering the batch into decode and prefill requests.
    # If > 1, the batch will be reordered such that requests with
    # query length <= threshold are classified as decode requests.
    # Use `supports_uniform_spec_as_decode` (above) to set this automatically
    # when speculative decoding is enabled.
477
    reorder_batch_threshold: int = 1
478

479
    @staticmethod
480
    def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int:
481
482
483
484
485
486
        scheduler_config = vllm_config.scheduler_config
        cache_config = vllm_config.cache_config
        model_config = vllm_config.model_config

        chunked_prefill_workspace_size = min(
            # Try for 8 full length request or at least 4 pages per-request
487
488
489
490
            max(
                8 * model_config.max_model_len,
                4 * scheduler_config.max_num_seqs * cache_config.block_size,
            ),
491
492
493
494
495
496
497
498
            # For long-context models try not to over-allocate limiting
            # kv-cache space, limiting it to 64k tokens,
            # which would result in the workspace being:
            #   2*(576)*(64*1024) = 144mb
            # (assuming 576 MLA head dim, and fp16)
            # which would result in up-projected context being
            #   2*(192*128)*(64*1024) = 3gb
            # (assuming 192 QK head dim, 128 heads, and fp16)
499
500
            64 * 1024,
        )
501
502
503
504

        # Enforce that we enough for at least 1 page per request
        chunked_prefill_workspace_size = max(
            chunked_prefill_workspace_size,
505
506
            scheduler_config.max_num_seqs * cache_config.block_size,
        )
507
508
509

        return chunked_prefill_workspace_size

510
511
512
513
514
515
    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
516
        metadata_cls: type[M] | None = None,
517
518
519
520
    ):
        self.metadata_cls = (
            metadata_cls if metadata_cls is not None else MLACommonMetadata
        )
521
522
523
524
        self.kv_cache_spec = kv_cache_spec
        scheduler_config = vllm_config.scheduler_config
        self.model_config = vllm_config.model_config
        parallel_config = vllm_config.parallel_config
525
        self.compilation_config = vllm_config.compilation_config
526
        self.vllm_config = vllm_config
527
528
        self.device = device

529
        self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
530
        self.mla_dims = get_mla_dims(self.model_config)
531
        self.aot_schedule = current_platform.is_cuda()
532
533
534
535
536
537
538
        try:
            self.dcp_world_size = get_dcp_group().world_size
            self.dcp_rank = get_dcp_group().rank_in_group
        except AssertionError:
            # DCP might not be initialized in testing
            self.dcp_world_size = 1
            self.dcp_rank = 0
539

540
        # Don't try to access the runner on AMD
541
        if self.aot_schedule:
542
            self.page_size = self.kv_cache_spec.block_size
543

544
        self.chunked_prefill_workspace_size = (
545
            self.determine_chunked_prefill_workspace_size(vllm_config)
546
        )
547

548
549
550
551
552
        if self.dcp_world_size > 1:
            # Note(hc): The local kvcache is incomplete when DCP is triggered,
            # an additional kvcache allgather across the DCP group is therefore
            # required, so the workspace has to be enlarged by 1/DCP relative
            # to the original TP allocation.
553
            assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
554
            self.chunked_prefill_workspace = torch.empty(
555
556
557
558
559
                (
                    self.chunked_prefill_workspace_size
                    + self.chunked_prefill_workspace_size // self.dcp_world_size,
                    self.model_config.get_head_size(),
                ),
560
561
562
563
564
                dtype=self.model_config.dtype,
                device=device,
            )
        else:
            self.chunked_prefill_workspace = torch.empty(
565
566
567
568
                (
                    self.chunked_prefill_workspace_size,
                    self.model_config.get_head_size(),
                ),
569
570
571
                dtype=self.model_config.dtype,
                device=device,
            )
572
573

        self._use_cudnn_prefill = use_cudnn_prefill()
574
        self._use_fi_prefill = use_flashinfer_prefill()
575
576
        self.prefill_metadata_cls = (
            FlashInferPrefillMetadata
577
578
579
580
581
            if self._use_fi_prefill
            else CudnnPrefillMetadata
            if self._use_cudnn_prefill
            else MLACommonPrefillMetadata
        )
582
583
584

        if self._use_fi_prefill:
            self._workspace_buffer = torch.empty(
585
586
                FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device
            )
587

588
            self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
589
            self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
590
591

            self._global_hyperparameters = infer_global_hyperparameters(
592
593
                get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl)
            )
594

595
596
597
598
        if self._use_cudnn_prefill:
            self.cudnn_workspace = torch.empty(
                CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,
                dtype=torch.int8,
599
                device=device,
600
601
            )

602
603
604
605
606
        supports_spec_as_decode = self.supports_uniform_spec_as_decode
        self._init_reorder_batch_threshold(
            self.reorder_batch_threshold, supports_spec_as_decode
        )

607
608
609
610
611
612
613
614
615
616
    def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):
        qo_indptr = prefill.query_start_loc

        has_context = False
        if prefill.chunked_context is not None:
            chunked_context = prefill.chunked_context
            has_context = True

        if self._fi_prefill_main is None:
            self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
617
618
                self._workspace_buffer, "NHD", backend="cutlass"
            )
619
620
621
622
623
624
625
626

        if has_context:
            num_chunks = chunked_context.cu_seq_lens.shape[0]
            # Allocate more prefill chunk wrappers if needed
            if len(self._fi_prefill_chunks) < num_chunks:
                for _ in range(len(self._fi_prefill_chunks), num_chunks):
                    self._fi_prefill_chunks.append(
                        BatchPrefillWithRaggedKVCacheWrapper(
627
628
629
                            self._workspace_buffer, "NHD", backend="cutlass"
                        )
                    )
630
631
632
            assert num_chunks <= len(self._fi_prefill_chunks)

        # In MLA, the non-latent num_qo_heads == num_kv_heads
633
        num_qo_heads = self.num_heads
634
635
636
637
638
639
        num_kv_heads = num_qo_heads

        # Sanity: Verify that num_kv_heads == 1 since it is latent space
        assert self.kv_cache_spec.num_kv_heads == 1

        # Get non-latent head_dim_qk and head_dim_vo
640
        head_dim_qk = self.mla_dims.qk_nope_head_dim + self.mla_dims.qk_rope_head_dim
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        head_dim_vo = self.mla_dims.v_head_dim

        # For main run, qo_indptr == kv_indptr
        kv_indptr = qo_indptr.clone()

        # Prepare main prefill
        self._fi_prefill_main.plan(
            qo_indptr=qo_indptr,
            kv_indptr=kv_indptr,
            num_qo_heads=num_qo_heads,
            num_kv_heads=num_kv_heads,
            head_dim_qk=head_dim_qk,
            head_dim_vo=head_dim_vo,
            causal=True,  # This is main run
            sm_scale=self._global_hyperparameters.sm_scale,
            window_left=self._global_hyperparameters.window_left,
            logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
658
            q_data_type=self.model_config.dtype,
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
        )

        # Prepare context prefills
        if has_context:
            for i in range(num_chunks):
                kv_indptr_chunk = chunked_context.cu_seq_lens[i]

                self._fi_prefill_chunks[i].plan(
                    qo_indptr=qo_indptr,
                    kv_indptr=kv_indptr_chunk,
                    num_qo_heads=num_qo_heads,
                    num_kv_heads=num_kv_heads,
                    head_dim_qk=head_dim_qk,
                    head_dim_vo=head_dim_vo,
                    causal=False,  # This is context run
                    sm_scale=self._global_hyperparameters.sm_scale,
                    window_left=self._global_hyperparameters.window_left,
676
                    logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
677
                    q_data_type=self.model_config.dtype,
678
679
680
681
682
                )

        prefill.prefill_main = self._fi_prefill_main
        prefill.prefill_chunks = self._fi_prefill_chunks

683
684
685
686
687
688
689
690
    def _build_decode(
        self,
        block_table_tensor: torch.Tensor,
        seq_lens_cpu: torch.Tensor,
        seq_lens_device: torch.Tensor,
        query_start_loc_cpu: torch.Tensor,
        query_start_loc_device: torch.Tensor,
        num_decode_tokens: int,
691
        dcp_tot_seq_lens_device: torch.Tensor | None,
692
    ) -> MLACommonDecodeMetadata:
693
        return MLACommonDecodeMetadata(
694
            block_table=block_table_tensor,
695
            seq_lens=seq_lens_device,
696
            dcp_tot_seq_lens=dcp_tot_seq_lens_device,
697
698
        )

699
    def build_for_cudagraph_capture(
700
701
        self, common_attn_metadata: CommonAttentionMetadata
    ) -> M:
702
703
704
705
706
        """
        This method builds the metadata for full cudagraph capture.
        Currently, only decode is supported for full cudagraphs with MLA.
        """
        m = common_attn_metadata
707
708
        assert m.num_reqs <= (m.num_actual_tokens * self.reorder_batch_threshold), (
            "MLA only supports decode-only full CUDAGraph capture. "
709
            "Make sure all cudagraph capture sizes <= max_num_seq."
710
        )
711

712
        assert m.max_query_len <= self.reorder_batch_threshold  # decode only
713
714
715

        return self.build(0, m)

716
717
718
719
720
721
    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> M:
722
        num_reqs = common_attn_metadata.num_reqs
723
        num_tokens = common_attn_metadata.num_actual_tokens
724
        max_query_len = common_attn_metadata.max_query_len
725
        max_seq_len = common_attn_metadata.max_seq_len
726

Simon Mo's avatar
Simon Mo committed
727
728
729
        # Note(simon): be careful about the CPU <> GPU memory movement in this
        # function. We should avoid GPU -> CPU sync as much as possible because
        # it blocks on all previous kernels.
730
731
732
        device = self.device
        block_table_tensor = common_attn_metadata.block_table_tensor
        slot_mapping = common_attn_metadata.slot_mapping
733

734
        query_start_loc = common_attn_metadata.query_start_loc
735
        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
736
        seq_lens = common_attn_metadata.seq_lens
737
        seq_lens_cpu = common_attn_metadata.seq_lens_cpu
738
        dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
Simon Mo's avatar
Simon Mo committed
739

740
741
        query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

742
        num_computed_tokens_cpu = common_attn_metadata.seq_lens_cpu - query_seq_lens_cpu
743

744
745
        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
            split_decodes_and_prefills(
746
747
748
                common_attn_metadata,
                decode_threshold=self.reorder_batch_threshold,
                require_uniform=self.supports_uniform_spec_as_decode,
749
750
            )
        )
751

752
753
        # Note(hc): update seq_lens of decode reqs under DCP.
        if self.dcp_world_size > 1:
754
755
756
757
            assert dcp_local_seq_lens is not None
            dcp_local_seq_lens[:num_decodes] = seq_lens[
                :num_decodes
            ] // self.dcp_world_size + (
758
759
                self.dcp_rank <= (seq_lens[:num_decodes] - 1) % self.dcp_world_size
            )
760

761
762
763
        assert num_decodes + num_prefills == num_reqs
        assert num_decode_tokens + num_prefill_tokens == num_tokens

764
        prefill_metadata = None
765
766
        if num_prefills > 0:
            reqs_start = num_decodes  # prefill_start
767

768
            context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
769
            # Note(hc): The context lengths in the perspective of dcp rank0.
770
771
772
            cp_context_lens_cpu = torch.ceil(
                context_lens_cpu.float() / self.dcp_world_size
            ).int()
773
            origin_context_lens = context_lens_cpu.tolist()
Simon Mo's avatar
Simon Mo committed
774
775
            max_context_len_cpu = context_lens_cpu.max().item()
            num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
776
777
778
            prefill_query_start_loc = (
                query_start_loc[reqs_start:] - query_start_loc[reqs_start]
            )
779
780

            chunked_context_metadata = None
781
            if max_context_len_cpu > 0:
782
783
784
785
786
787
788
789
                # NOTE: it is recommend you read the `Chunked Prefill` section
                # in the comment at the top of the file before trying to
                # understand the following code

                # currently we allocate an equal amount of workspace for each
                # prefill in the batch, we could probably use a more advanced
                # algorithm here and allocate more workspace to prefills with
                # longer context lengths
790
791
792
                max_context_chunk = (
                    self.chunked_prefill_workspace_size // num_prefills_with_context_cpu
                )
793

794
795
                if self.aot_schedule:
                    # align max_context_chunk to page_size by rounding down,
796
797
798
                    # currently the `gather_and_maybe_dequant_cache` kernel
                    # cannot handle `context_chunk_starts` that are not aligned
                    # to page_size
799
                    max_context_chunk = round_down(max_context_chunk, self.page_size)
800
801

                assert max_context_chunk > 0
Simon Mo's avatar
Simon Mo committed
802
                num_chunks = cdiv(max_context_len_cpu, max_context_chunk)
803
804
805
806
807

                # if `max_context_chunk = 256`, `num_chunks = 3`, and
                #   `num_prefills_with_context = 4`, create a tensor that looks
                # like
                #  [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
Simon Mo's avatar
Simon Mo committed
808
809
                # Note(simon): this is done in CPU because of downstream's
                # of `to_list`.
810
811
812
813
                chunk_starts = (
                    torch.arange(num_chunks, dtype=torch.int32)
                    .unsqueeze(1)
                    .expand(-1, num_prefills)
814
                    * max_context_chunk
815
816
817
818
                )
                chunk_ends = torch.min(
                    context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
                )
819
                chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
Simon Mo's avatar
Simon Mo committed
820

821
822
823
824
825
826
                cu_seq_lens_cpu = torch.zeros(
                    num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
                )
                torch.cumsum(
                    chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
                )
827

828
829
830
831
832
833
834
                if self.dcp_world_size > 1:
                    # Note(hc): The above max_context_chunk already enforces
                    # block_size alignment, DCP just need the block_size can
                    # be divisible by dcp_world_size, because DCP use
                    # cp_gather_cache which not require `cp_chunk_starts`
                    # aligned to page_size.
                    assert max_context_chunk % self.dcp_world_size == 0
835
836
837
838
839
                    cp_max_context_chunk = max_context_chunk // self.dcp_world_size
                    cp_chunk_starts = (
                        torch.arange(num_chunks, dtype=torch.int32)
                        .unsqueeze(1)
                        .expand(-1, num_prefills)
840
                        * cp_max_context_chunk
841
                    )
842
843
                    cp_chunk_ends = torch.min(
                        cp_context_lens_cpu.unsqueeze(0),
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
                        cp_chunk_starts + cp_max_context_chunk,
                    )
                    cp_chunk_seq_lens = (cp_chunk_ends - cp_chunk_starts).clamp(min=0)

                    cp_cu_seq_lens_cpu = torch.zeros(
                        num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
                    )
                    torch.cumsum(
                        cp_chunk_seq_lens,
                        dim=1,
                        out=cp_cu_seq_lens_cpu[:, 1:],
                        dtype=torch.int32,
                    )

                chunked_context_metadata_cls = (
                    CudnnPrefillMetadata.ChunkedContextMetadata
                    if self._use_cudnn_prefill
                    else MLACommonPrefillMetadata.ChunkedContextMetadata
                )
863
                if self.dcp_world_size > 1:
864
865
                    chunked_context_metadata = chunked_context_metadata_cls(
                        cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
866
867
868
869
870
871
872
                        starts=cp_chunk_starts.to(device, non_blocking=True),
                        seq_tot=cp_chunk_seq_lens.sum(dim=1).tolist(),
                        max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
                        seq_lens=chunk_seq_lens,
                        workspace=self.chunked_prefill_workspace,
                        cp_chunk_seq_lens=cp_chunk_seq_lens.tolist(),
                        origin_context_lens=origin_context_lens,
873
                        cp_cu_seq_lens=cp_cu_seq_lens_cpu.to(device, non_blocking=True),
874
875
876
877
                        chunk_size=max_context_chunk,
                        cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
                    )
                else:
878
879
                    chunked_context_metadata = chunked_context_metadata_cls(
                        cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
880
881
882
883
884
885
                        starts=chunk_starts.to(device, non_blocking=True),
                        seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
                        max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
                        seq_lens=chunk_seq_lens,
                        workspace=self.chunked_prefill_workspace,
                    )
886

887
888
889
                if self._use_cudnn_prefill:
                    chunked_context_metadata.seq_lens = chunk_seq_lens

890
891
892
893
                assert (
                    max(chunked_context_metadata.max_seq_lens)
                    <= self.chunked_prefill_workspace_size
                )
894

895
            prefill_metadata = self.prefill_metadata_cls(
896
                block_table=block_table_tensor[reqs_start:, ...],
897
                query_start_loc=prefill_query_start_loc,
Simon Mo's avatar
Simon Mo committed
898
                max_query_len=max_query_len,
899
900
901
                chunked_context=chunked_context_metadata,
            )

902
903
            if self._use_cudnn_prefill:
                assert isinstance(prefill_metadata, CudnnPrefillMetadata)
904
905
906
                prefill_metadata.query_seq_lens = (
                    prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]
                )
907
908
                prefill_metadata.cudnn_workspace = self.cudnn_workspace

909
        decode_metadata = None
910
        if num_decodes > 0:
911
            decode_metadata = self._build_decode(
912
                block_table_tensor=block_table_tensor[:num_decodes, ...],
913
                seq_lens_cpu=seq_lens_cpu[:num_decodes],
914
915
916
                seq_lens_device=dcp_local_seq_lens[:num_decodes]
                if self.dcp_world_size > 1 and dcp_local_seq_lens is not None
                else seq_lens[:num_decodes],
917
918
                query_start_loc_cpu=query_start_loc_cpu[: num_decodes + 1],
                query_start_loc_device=query_start_loc[: num_decodes + 1],
919
                num_decode_tokens=num_decode_tokens,
920
921
922
                dcp_tot_seq_lens_device=seq_lens[:num_decodes]
                if self.dcp_world_size > 1
                else None,
923
924
            )

925
        attn_metadata = self.metadata_cls(
926
927
            num_reqs=common_attn_metadata.num_reqs,
            max_query_len=common_attn_metadata.max_query_len,
928
            max_seq_len=max_seq_len,
929
            num_actual_tokens=num_tokens,
930
931
            query_start_loc=query_start_loc,
            slot_mapping=slot_mapping,
932
            head_dim=self.model_config.get_head_size(),
933
            # MLACommonMetadata Chunk prefill specific
934
935
936
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
937
938
            prefill=prefill_metadata,
            decode=decode_metadata,
939
940
        )

941
        if self._use_fi_prefill and num_prefills > 0:
942
943
944
945
946
            assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)
            self._build_fi_prefill_wrappers(attn_metadata.prefill)

        return attn_metadata

947

948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
def reorg_kvcache(
    allgatered_kv_c_normed: torch.Tensor,
    allgatered_k_pe: torch.Tensor,
    cp_chunk_seq_lens_lst: list[int],
    origin_context_lens: list[int],
    cp_world_size: int,
    sum_seq_len: int,
    max_seq_len: int,
    chunk_size: int,
    chunk_idx: int,
    toks: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    reorg kvcache after cp local gather to tp layout for attn kernel.

    Args:
        cp_chunk_seq_lens_lst: chunk context lengths under CP.
        origin_context_lens: origin full context lengths under CP.
        cp_world_size: CP size.
        sum_seq_len: the sum of cp_chunk_seq_lens_lst.
        max_seq_len: the max value of cp_chunk_seq_lens_lst.
        chunk_size: equals to max_context_chunk from
            chunked_context_metadata building.
        chunk_idx: chunk idx of chunked_prefill.
        toks: the number of tokens for local gather cache.
    """
    kv_c_segments = []
    k_pe_segments = []
    src_token_idx = 0
    max_seq_len_check = 0
978
979
980
    for cp_chunk_seq_len, origin_context_len in zip(
        cp_chunk_seq_lens_lst, origin_context_lens
    ):
981
982
983
        chunk_context_len = chunk_size
        if cp_chunk_seq_len != 0:
            chunk_context_len = min(
984
985
                chunk_context_len, origin_context_len - chunk_size * chunk_idx
            )
986
987
988
989
990
991
992
993
        cp_target_rank = (chunk_context_len - 1) % cp_world_size
        cur_seq_len = 0
        for rank in range(cp_world_size):
            if rank > cp_target_rank and cp_chunk_seq_len:
                real_cp_chunk_seq_len = cp_chunk_seq_len - 1
            else:
                real_cp_chunk_seq_len = cp_chunk_seq_len
            if real_cp_chunk_seq_len:
994
995
996
997
998
999
1000
1001
1002
1003
                kv_c_segment = allgatered_kv_c_normed[
                    rank * toks + src_token_idx : rank * toks
                    + src_token_idx
                    + real_cp_chunk_seq_len
                ]
                k_pe_segment = allgatered_k_pe[
                    rank * toks + src_token_idx : rank * toks
                    + src_token_idx
                    + real_cp_chunk_seq_len
                ]
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
                kv_c_segments.append(kv_c_segment)
                k_pe_segments.append(k_pe_segment)
                cur_seq_len += real_cp_chunk_seq_len
        max_seq_len_check = max(max_seq_len_check, cur_seq_len)
        src_token_idx += cp_chunk_seq_len
    reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
    reorganized_k_pe = torch.cat(k_pe_segments, dim=0)
    assert reorganized_kv_c_normed.shape[0] == sum_seq_len
    assert reorganized_k_pe.shape[0] == sum_seq_len
    assert max_seq_len_check == max_seq_len
    return reorganized_kv_c_normed, reorganized_k_pe


1017
1018
1019
# TODO(Lucas): rename MLACommonBaseImpl -> MLACommonImpl,
# and MLACommonImpl -> MLACommonDenseImpl or somthing like that
class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
    """
    NOTE: Please read the comment at the top of the file before trying to
    understand this class
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
1031
1032
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
1033
        kv_cache_dtype: str,
1034
        logits_soft_cap: float | None,
1035
        attn_type: str,
1036
        kv_sharing_target_layer_name: str | None,
1037
        # MLA Specific Arguments
1038
        q_lora_rank: int | None,
1039
1040
1041
1042
1043
1044
        kv_lora_rank: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        qk_head_dim: int,
        v_head_dim: int,
        kv_b_proj: ColumnParallelLinear,
1045
        indexer=None,
1046
        q_pad_num_heads: int | None = None,
1047
    ) -> None:
1048
1049
1050
        if kv_sharing_target_layer_name is not None:
            raise NotImplementedError("KV sharing is not supported for MLA")

1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        self.kv_cache_dtype = kv_cache_dtype

        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_head_dim
        self.v_head_dim = v_head_dim
        self.kv_b_proj = kv_b_proj
1064
        self.indexer = indexer
1065
        self.q_pad_num_heads = q_pad_num_heads
1066

1067
1068
1069
1070
1071
1072
1073
    def process_weights_after_loading(self, act_dtype: torch.dtype):
        def get_layer_weight(layer):
            WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
            for attr in WEIGHT_NAMES:
                if hasattr(layer, attr):
                    return getattr(layer, attr)
            raise AttributeError(
1074
1075
                f"Layer '{layer}' has no recognized weight attribute: {WEIGHT_NAMES}."
            )
1076
1077
1078
1079

        def get_and_maybe_dequant_weights(layer: LinearBase):
            if not isinstance(layer.quant_method, UnquantizedLinearMethod):
                # NOTE: This should only be used offline, since it's O(N^3)
1080
1081
1082
1083
1084
1085
                eye = torch.eye(
                    layer.input_size_per_partition,
                    dtype=act_dtype,
                    device=get_layer_weight(layer).device,
                )
                dequant_weights = layer.quant_method.apply(layer, eye, bias=None)
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
                del eye
                # standardize to (output, input)
                return dequant_weights.T
            return layer.weight

        # we currently do not have quantized bmm's which are needed for
        # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
        # the bmm's in 16-bit, the extra memory overhead of this is fairly low
        kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
        assert kv_b_proj_weight.shape == (
            self.kv_lora_rank,
1097
1098
1099
1100
1101
1102
1103
1104
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
        ), (
            f"{kv_b_proj_weight.shape=}, "
            f"{self.kv_lora_rank=}, "
            f"{self.num_heads=}, "
            f"{self.qk_nope_head_dim=}, "
            f"{self.v_head_dim=}"
        )
1105
1106
1107
1108
1109
1110
1111
        kv_b_proj_weight = kv_b_proj_weight.view(
            self.kv_lora_rank,
            self.num_heads,
            self.qk_nope_head_dim + self.v_head_dim,
        )

        W_UK, W_UV = kv_b_proj_weight.split(
1112
1113
            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )
1114
1115
1116
1117
1118

        if is_rocm_aiter_fp8bmm_enabled():
            W_K = W_UK.transpose(0, 1)  # 16 512 128
            W_V = W_UV.permute(1, 2, 0)  # 16 128 512
            self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
1119
1120
                W_K, dtype=current_platform.fp8_dtype()
            )
1121
            self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
1122
1123
                W_V, dtype=current_platform.fp8_dtype()
            )
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138

            # The kernel operates on non-padded inputs. Hence, pre-compiling
            # triton kernel to avoid runtime compilation for unseen batch sizes
            # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
            # On DS-R1, this step adds roughly 50s to the model loading time.
            max_batch_size = 1024  # [ToDo] Find the optimal upper limit
            pre_compilation_list = list(range(1, max_batch_size + 1))
            if is_global_first_rank():
                pre_compilation_list = tqdm(
                    pre_compilation_list,
                    desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
                    total=max_batch_size,
                )

            for m in pre_compilation_list:
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
                x = torch.empty(
                    (self.W_K.shape[0], m, self.W_K.shape[2]),
                    dtype=torch.bfloat16,
                    device=self.W_K.device,
                )
                aiter_triton_fp8_bmm(
                    x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
                )

                x = torch.empty(
                    (self.W_V.shape[0], m, self.W_V.shape[2]),
                    dtype=torch.bfloat16,
                    device=self.W_V.device,
                )
                aiter_triton_fp8_bmm(
                    x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
                )
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
        else:
            # Convert from (L, N, V) to (N, L, V)
            self.W_UV = W_UV.transpose(0, 1)
            # Convert from (L, N, P) to (N, P, L)
            self.W_UK_T = W_UK.permute(1, 2, 0)

    def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
        # Convert from (B, N, L) to (N, B, L)
        x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
        if is_rocm_aiter_fp8bmm_enabled():
            # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1167
1168
1169
            x = aiter_triton_fp8_bmm(
                x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
            )
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
            # Convert from (B, N, V) to (B, N * V)
            x = x.reshape(-1, self.num_heads * self.v_head_dim)
            # Copy result
            out.copy_(x)
        else:
            # Convert from (B, N * V) to (N, B, V)
            out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)

            # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
            torch.bmm(x, self.W_UV, out=out)  # Reuse "out" to make it "hot"

            # Convert from (N, B, V) to (B, N * V)
1182
            out_new = out.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198

            # Adjust output buffer shape back to the original (B, N * V)
            N, B, V = out.shape
            out.resize_((B, N * V))
            out.copy_(out_new)  # Copy result


class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
    """
    NOTE: Please read the comment at the top of the file before trying to
    understand this class
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

1199
1200
1201
1202
1203
        if use_flashinfer_prefill():
            logger.debug_once("Using FlashInfer prefill for MLA")
            self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
            self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
            self._pad_v = False
1204
1205
        elif use_cudnn_prefill():
            logger.debug_once("Using CUDNN prefill for MLA")
1206
            self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn
1207
1208
            self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn
            self._pad_v = False
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
        else:  # Use FlashAttention
            logger.debug_once("Using FlashAttention prefill for MLA")
            self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa
            self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa

            # Handle the differences between the flash_attn_varlen from
            # flash_attn and the one from vllm_flash_attn. The former is used on
            # RoCM and the latter has an additional parameter to control
            # FA2 vs FA3
            self.flash_attn_varlen_func = flash_attn_varlen_func
            self.vllm_flash_attn_version = get_flash_attn_version()
            if self.vllm_flash_attn_version is not None:
1221
1222
1223
                self.flash_attn_varlen_func = functools.partial(
                    flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
                )
1224
1225
1226
1227
1228
1229
1230

            # For MLA the v head dim is smaller than qk head dim so we pad out
            # v with 0s to match the qk head dim for attention backends that do
            # not support different headdims
            # We don't need to pad V if we are on a hopper system with FA3
            self._pad_v = self.vllm_flash_attn_version is None or not (
                self.vllm_flash_attn_version == 3
1231
1232
                and current_platform.get_device_capability()[0] == 9
            )
1233

1234
        self.dcp_world_size: int | None = None
1235

1236
        self.chunked_prefill_workspace_size = (
1237
            MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(
1238
1239
1240
1241
1242
1243
1244
                get_current_vllm_config()
            )
        )

    def _flash_attn_varlen_diff_headdims(
        self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
    ):
1245
1246
1247
        maybe_padded_v = v
        if self._pad_v:
            maybe_padded_v = torch.nn.functional.pad(
1248
1249
                v, [0, q.shape[-1] - v.shape[-1]], value=0
            )
1250

1251
1252
1253
1254
1255
1256
1257
        if is_vllm_fa:
            kwargs["return_softmax_lse"] = return_softmax_lse
        else:
            # ROCm leverages the upstream flash_attn, which takes a parameter
            # called "return_attn_probs" instead of return_softmax_lse
            kwargs["return_attn_probs"] = return_softmax_lse

1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
        attn_out = self.flash_attn_varlen_func(
            q=q,
            k=k,
            v=maybe_padded_v,
            softmax_scale=softmax_scale,
            **kwargs,
        )

        # Unpack the output if there is multiple results
        lse = None
        if isinstance(attn_out, tuple):
            attn_out, lse = attn_out[0], attn_out[1]

        # Remain consistent with old `flash_attn_varlen_func` where there
        # is only one output tensor if `return_softmax_lse` is False.
        if return_softmax_lse:
            return attn_out, lse
        return attn_out

1277
1278
1279
    def _run_prefill_new_tokens_fa(
        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
    ):
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
        return self._flash_attn_varlen_diff_headdims(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=prefill.query_start_loc,
            cu_seqlens_k=prefill.query_start_loc,
            max_seqlen_q=prefill.max_query_len,
            max_seqlen_k=prefill.max_query_len,
            softmax_scale=self.scale,
            causal=True,
            return_softmax_lse=return_softmax_lse,
        )

1293
1294
1295
    def _run_prefill_new_tokens_fi(
        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
    ):
1296
1297
        assert isinstance(prefill, FlashInferPrefillMetadata)
        assert prefill.prefill_main is not None
1298
        ret = prefill.prefill_main.run(
1299
1300
1301
1302
1303
1304
            q=q,
            k=k,
            v=v,
            return_lse=return_softmax_lse,
        )

1305
1306
1307
1308
1309
        if isinstance(ret, tuple):
            # Convert from (q_len, num_heads) to (num_heads, q_len)
            return ret[0], ret[1].transpose(0, 1).contiguous()
        return ret

1310
1311
1312
    def _run_prefill_new_tokens_cudnn(
        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
    ):
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
        assert isinstance(prefill, CudnnPrefillMetadata)
        assert prefill.query_seq_lens is not None
        output, lse = cudnn_batch_prefill_with_kv_cache(
            q=q,
            k_cache=k,
            v_cache=v,
            scale=self.scale,
            workspace_buffer=prefill.cudnn_workspace,
            max_token_per_sequence=prefill.max_query_len,
            max_sequence_kv=prefill.max_query_len,
            actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),
            actual_seq_lens_kv=prefill.query_seq_lens.view(-1, 1, 1, 1),
            causal=True,
1326
1327
1328
1329
            # Do not support False for now
            return_lse=True,
            # Indicates actual_seq_lens are on GPU or CPU.
            is_cuda_graph_compatible=True,
1330
1331
1332
1333
1334
        )
        if return_softmax_lse:
            return output, lse
        return output

1335
1336
1337
    def _run_prefill_context_chunk_fa(
        self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
    ):
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
        assert prefill.chunked_context is not None
        return self._flash_attn_varlen_diff_headdims(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=prefill.query_start_loc,
            cu_seqlens_k=prefill.chunked_context.cu_seq_lens[chunk_idx],
            max_seqlen_q=prefill.max_query_len,
            max_seqlen_k=prefill.chunked_context.max_seq_lens[chunk_idx],
            softmax_scale=self.scale,
            causal=False,  # Context is unmasked
            return_softmax_lse=True,
        )

1352
1353
1354
    def _run_prefill_context_chunk_fi(
        self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
    ):
1355
        assert isinstance(prefill, FlashInferPrefillMetadata)
1356
        attn_out, lse = prefill.prefill_chunks[chunk_idx].run(
1357
1358
1359
1360
1361
            q=q,
            k=k,
            v=v,
            return_lse=True,
        )
1362
1363
        # Convert from (q_len, num_heads) to (num_heads, q_len)
        return attn_out, lse.transpose(0, 1).contiguous()
1364

1365
1366
1367
    def _run_prefill_context_chunk_cudnn(
        self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
    ):
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
        assert isinstance(prefill, CudnnPrefillMetadata)
        assert prefill.chunked_context is not None
        assert prefill.chunked_context.seq_lens[chunk_idx] is not None
        assert prefill.query_seq_lens is not None
        return cudnn_batch_prefill_with_kv_cache(
            q=q,
            k_cache=k,
            v_cache=v,
            scale=self.scale,
            workspace_buffer=prefill.cudnn_workspace,
            max_token_per_sequence=prefill.max_query_len,
            max_sequence_kv=prefill.chunked_context.max_seq_lens[chunk_idx],
            actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),
1381
1382
1383
            actual_seq_lens_kv=prefill.chunked_context.seq_lens[chunk_idx].view(
                -1, 1, 1, 1
            ),
1384
1385
            causal=False,
            return_lse=True,
1386
1387
            # Indicates actual_seq_lens are on GPU or CPU.
            is_cuda_graph_compatible=True,
1388
1389
        )

1390
    def process_weights_after_loading(self, act_dtype: torch.dtype):
1391
        def get_layer_weight(layer):
1392
1393
1394
1395
1396
            WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
            for attr in WEIGHT_NAMES:
                if hasattr(layer, attr):
                    return getattr(layer, attr)
            raise AttributeError(
1397
1398
                f"Layer '{layer}' has no recognized weight attribute: {WEIGHT_NAMES}."
            )
1399
1400
1401
1402

        def get_and_maybe_dequant_weights(layer: LinearBase):
            if not isinstance(layer.quant_method, UnquantizedLinearMethod):
                # NOTE: This should only be used offline, since it's O(N^3)
1403
1404
1405
1406
1407
1408
                eye = torch.eye(
                    layer.input_size_per_partition,
                    dtype=act_dtype,
                    device=get_layer_weight(layer).device,
                )
                dequant_weights = layer.quant_method.apply(layer, eye, bias=None)
1409
1410
1411
1412
1413
                del eye
                # standardize to (output, input)
                return dequant_weights.T
            return layer.weight

1414
        # we currently do not have quantized bmm's which are needed for
1415
        # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1416
        # the bmm's in 16-bit, the extra memory overhead of this is fairly low
1417
1418
1419
        kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
        assert kv_b_proj_weight.shape == (
            self.kv_lora_rank,
1420
1421
1422
1423
1424
1425
1426
1427
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
        ), (
            f"{kv_b_proj_weight.shape=}, "
            f"{self.kv_lora_rank=}, "
            f"{self.num_heads=}, "
            f"{self.qk_nope_head_dim=}, "
            f"{self.v_head_dim=}"
        )
1428
1429
1430
1431
1432
1433
1434
        kv_b_proj_weight = kv_b_proj_weight.view(
            self.kv_lora_rank,
            self.num_heads,
            self.qk_nope_head_dim + self.v_head_dim,
        )

        W_UK, W_UV = kv_b_proj_weight.split(
1435
1436
            [self.qk_nope_head_dim, self.v_head_dim], dim=-1
        )
1437

1438
1439
1440
1441
        if is_rocm_aiter_fp8bmm_enabled():
            W_K = W_UK.transpose(0, 1)  # 16 512 128
            W_V = W_UV.permute(1, 2, 0)  # 16 128 512
            self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(
1442
1443
                W_K, dtype=current_platform.fp8_dtype()
            )
1444
            self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(
1445
1446
                W_V, dtype=current_platform.fp8_dtype()
            )
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461

            # The kernel operates on non-padded inputs. Hence, pre-compiling
            # triton kernel to avoid runtime compilation for unseen batch sizes
            # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
            # On DS-R1, this step adds roughly 50s to the model loading time.
            max_batch_size = 1024  # [ToDo] Find the optimal upper limit
            pre_compilation_list = list(range(1, max_batch_size + 1))
            if is_global_first_rank():
                pre_compilation_list = tqdm(
                    pre_compilation_list,
                    desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
                    total=max_batch_size,
                )

            for m in pre_compilation_list:
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
                x = torch.empty(
                    (self.W_K.shape[0], m, self.W_K.shape[2]),
                    dtype=torch.bfloat16,
                    device=self.W_K.device,
                )
                aiter_triton_fp8_bmm(
                    x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True
                )

                x = torch.empty(
                    (self.W_V.shape[0], m, self.W_V.shape[2]),
                    dtype=torch.bfloat16,
                    device=self.W_V.device,
                )
                aiter_triton_fp8_bmm(
                    x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True
                )
1479
1480
1481
1482
1483
        else:
            # Convert from (L, N, V) to (N, L, V)
            self.W_UV = W_UV.transpose(0, 1)
            # Convert from (L, N, P) to (N, P, L)
            self.W_UK_T = W_UK.permute(1, 2, 0)
1484
1485
1486
1487
1488
1489

    def _compute_prefill_context(
        self,
        q: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
1490
        k_scale: torch.Tensor,
1491
    ):
1492
1493
1494
        assert attn_metadata.prefill is not None
        prefill_metadata = attn_metadata.prefill
        assert prefill_metadata.chunked_context is not None
1495
1496

        output = None
1497
1498
        iters = len(prefill_metadata.chunked_context.seq_tot)
        workspace = prefill_metadata.chunked_context.workspace
1499
1500

        for i in range(iters):
1501
            toks = prefill_metadata.chunked_context.seq_tot[i]
1502

1503
            ops.gather_and_maybe_dequant_cache(
1504
1505
                src_cache=kv_c_and_k_pe_cache,
                dst=workspace,
1506
1507
                block_table=prefill_metadata.block_table,
                cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
1508
                batch_size=attn_metadata.num_prefills,
1509
1510
                kv_cache_dtype=self.kv_cache_dtype,
                scale=k_scale,
1511
                seq_starts=prefill_metadata.chunked_context.starts[i],
1512
1513
            )

1514
1515
            kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]
            k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
1516

1517
1518
1519
1520
            kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
                -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
            )
            k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1521

1522
            k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1523

1524
1525
1526
            attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
                prefill=prefill_metadata,
                chunk_idx=i,
1527
1528
                q=q,
                k=k,
1529
                v=v,
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
            )

            if output is None:
                output = attn_output
                output_lse = attn_softmax_lse
            else:
                output_tmp = torch.empty_like(output)
                output_lse_tmp = torch.empty_like(output_lse)
                merge_attn_states(
                    output=output_tmp,
                    output_lse=output_lse_tmp,
                    prefix_output=output,
                    prefix_lse=output_lse,
                    suffix_output=attn_output,
                    suffix_lse=attn_softmax_lse,
                )
                output = output_tmp
                output_lse = output_lse_tmp

        return output, output_lse

1551
1552
1553
1554
1555
1556
1557
1558
    def _context_parallel_compute_prefill_context(
        self,
        q: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
        k_scale: torch.Tensor,
        dcp_world_size: int,
    ):
co63oc's avatar
co63oc committed
1559
        assert k_scale is None, "DCP not support scaled kvcache now."
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
        assert attn_metadata.prefill is not None
        prefill_metadata = attn_metadata.prefill
        assert prefill_metadata.chunked_context is not None
        assert prefill_metadata.chunked_context.cp_chunk_seq_lens is not None
        assert prefill_metadata.chunked_context.origin_context_lens is not None
        assert prefill_metadata.chunked_context.cp_cu_seq_lens is not None
        assert prefill_metadata.chunked_context.chunk_size is not None
        assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None

        output = None
        iters = len(prefill_metadata.chunked_context.seq_tot)
        workspace = prefill_metadata.chunked_context.workspace

        for i in range(iters):
            toks = prefill_metadata.chunked_context.seq_tot[i]
            ops.cp_gather_cache(
                src_cache=kv_c_and_k_pe_cache,
                dst=workspace,
                block_table=prefill_metadata.block_table,
                cu_seq_lens=prefill_metadata.chunked_context.cp_cu_seq_lens[i],
                batch_size=attn_metadata.num_prefills,
                seq_starts=prefill_metadata.chunked_context.starts[i],
            )
            # workspace
            # |------- N tokens --------|--------- N*dcp_size tokens ----------|
            # |<- use for loca_gather ->|<--------- use for allgather -------->|
            allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
1587
            assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0]
1588
1589
1590
            assert toks <= allgather_offset
            local_gathered_kvcache = workspace[:toks]
            cur_allgather_workspace = workspace[
1591
1592
                allgather_offset : allgather_offset * (1 + dcp_world_size)
            ]
1593
            assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
            cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
            cur_allgather_kvcache.copy_(
                get_dcp_group().all_gather(local_gathered_kvcache, dim=0)
            )
            assert (
                cur_allgather_kvcache.shape[-1]
                == self.kv_lora_rank + self.qk_rope_head_dim
            )
            allgatered_kv_c_normed, allgatered_k_pe = cur_allgather_kvcache.unsqueeze(
                1
            ).split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
1605
1606
1607
1608

            kv_c_normed, k_pe = reorg_kvcache(
                allgatered_kv_c_normed,
                allgatered_k_pe,
1609
1610
1611
1612
                cp_chunk_seq_lens_lst=prefill_metadata.chunked_context.cp_chunk_seq_lens[
                    i
                ],
                origin_context_lens=prefill_metadata.chunked_context.origin_context_lens,
1613
                cp_world_size=dcp_world_size,
1614
                sum_seq_len=prefill_metadata.chunked_context.cu_seq_lens_lst[i][-1],
1615
1616
1617
                max_seq_len=prefill_metadata.chunked_context.max_seq_lens[i],
                chunk_size=prefill_metadata.chunked_context.chunk_size,
                chunk_idx=i,
1618
1619
                toks=toks,
            )
1620

1621
1622
1623
1624
1625
            kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
                -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
            )
            k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
            k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653

            attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
                prefill=prefill_metadata,
                chunk_idx=i,
                q=q,
                k=k,
                v=v,
            )

            if output is None:
                output = attn_output
                output_lse = attn_softmax_lse
            else:
                output_tmp = torch.empty_like(output)
                output_lse_tmp = torch.empty_like(output_lse)
                merge_attn_states(
                    output=output_tmp,
                    output_lse=output_lse_tmp,
                    prefix_output=output,
                    prefix_lse=output_lse,
                    suffix_output=attn_output,
                    suffix_lse=attn_softmax_lse,
                )
                output = output_tmp
                output_lse = output_lse_tmp

        return output, output_lse

1654
1655
1656
1657
1658
1659
1660
    def _forward_prefill(
        self,
        q: torch.Tensor,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        kv_c_and_k_pe_cache: torch.Tensor,
        attn_metadata: MLACommonMetadata,
1661
        k_scale: torch.Tensor,
1662
    ) -> torch.Tensor:
1663
        # TODO (zyongye): Prefill function here
1664
        assert attn_metadata.prefill is not None
1665
        assert self.dcp_world_size is not None
1666
1667

        has_context = attn_metadata.prefill.chunked_context is not None
1668
1669
1670
1671
        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
            -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
        )
        k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1672
1673
1674

        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)

1675
1676
        output = self._run_prefill_new_tokens(
            prefill=attn_metadata.prefill,
1677
1678
            q=q,
            k=k,
1679
            v=v,
1680
1681
1682
1683
1684
            return_softmax_lse=has_context,
        )

        if has_context:
            suffix_output, suffix_lse = output
1685
            if self.dcp_world_size > 1:
1686
                context_output, context_lse = (
1687
                    self._context_parallel_compute_prefill_context(
1688
1689
1690
1691
1692
1693
1694
                        q,
                        kv_c_and_k_pe_cache,
                        attn_metadata,
                        k_scale=None,
                        dcp_world_size=self.dcp_world_size,
                    )
                )
1695
            else:
1696
1697
1698
                context_output, context_lse = self._compute_prefill_context(
                    q, kv_c_and_k_pe_cache, attn_metadata, k_scale
                )
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708

            output = torch.empty_like(suffix_output)
            merge_attn_states(
                output=output,
                prefix_output=context_output,
                prefix_lse=context_lse,
                suffix_output=suffix_output,
                suffix_lse=suffix_lse,
            )

1709
1710
        # unpad if necessary
        if self._pad_v:
1711
            output = output[..., : v.shape[-1]]
1712

1713
        return output.flatten(start_dim=-2)
1714
1715
1716
1717

    @abstractmethod
    def _forward_decode(
        self,
1718
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
1719
        kv_c_and_k_pe_cache: torch.Tensor,
1720
        attn_metadata: M,
1721
        layer: AttentionLayer,
1722
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
1723
1724
1725
1726
1727
        raise NotImplementedError

    def forward(
        self,
        layer: AttentionLayer,
1728
        q: torch.Tensor,
1729
1730
1731
        k_c_normed: torch.Tensor,  # key in unified attn
        k_pe: torch.Tensor,  # value in unified attn
        kv_cache: torch.Tensor,
1732
        attn_metadata: M,
1733
1734
1735
        output: torch.Tensor | None = None,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
1736
1737
1738
    ) -> torch.Tensor:
        assert output is not None, "Output tensor must be provided."

1739
        if output_scale is not None or output_block_scale is not None:
1740
            raise NotImplementedError(
1741
1742
                "fused output quantization is not yet supported for MLACommonImpl"
            )
1743

1744
        if attn_metadata is None:
1745
1746
1747
1748
            # During the profile run try to simulate to worse case output size
            # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
            # since this can be large
            _ = torch.empty(
1749
1750
1751
1752
1753
                (
                    self.chunked_prefill_workspace_size,
                    self.num_heads,
                    self.qk_nope_head_dim + self.v_head_dim,
                ),
1754
1755
1756
1757
                device=k_c_normed.device,
                dtype=k_c_normed.dtype,
            )

1758
1759
1760
1761
            # The zero fill is required when used with DP + EP
            # to ensure all ranks within a DP group compute the
            # same expert outputs.
            return output.fill_(0)
1762

1763
1764
1765
        if self.dcp_world_size is None:
            self.dcp_world_size = get_dcp_group().world_size

1766
1767
        fp8_attention = self.kv_cache_dtype.startswith("fp8")

1768
1769
1770
1771
1772
        num_actual_toks = attn_metadata.num_actual_tokens

        # Inputs and outputs may be padded for CUDA graphs
        output_padded = output
        output = output[:num_actual_toks, ...]
1773
        q = q[:num_actual_toks, ...]
1774
1775
1776
        k_c_normed = k_c_normed[:num_actual_toks, ...]
        k_pe = k_pe[:num_actual_toks, ...]

1777
1778
1779
1780
1781
        assert (
            attn_metadata.num_decodes is not None
            and attn_metadata.num_prefills is not None
            and attn_metadata.num_decode_tokens is not None
        )
1782
1783
1784
1785
1786

        has_decode = attn_metadata.num_decodes > 0
        has_prefill = attn_metadata.num_prefills > 0
        num_decode_tokens = attn_metadata.num_decode_tokens

1787
        decode_q = q[:num_decode_tokens]
1788

1789
        prefill_q = q[num_decode_tokens:]
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
        prefill_k_pe = k_pe[num_decode_tokens:]
        prefill_k_c_normed = k_c_normed[num_decode_tokens:]

        # write the latent and rope to kv cache
        if kv_cache.numel() > 0:
            ops.concat_and_cache_mla(
                k_c_normed,
                k_pe.squeeze(1),
                kv_cache,
                attn_metadata.slot_mapping.flatten(),
                kv_cache_dtype=self.kv_cache_dtype,
                scale=layer._k_scale,
            )

1804
1805
1806
        if fp8_attention:
            kv_cache = kv_cache.view(current_platform.fp8_dtype())

1807
1808
        if has_prefill:
            output[num_decode_tokens:] = self._forward_prefill(
1809
1810
1811
1812
1813
1814
1815
                prefill_q,
                prefill_k_c_normed,
                prefill_k_pe,
                kv_cache,
                attn_metadata,
                layer._k_scale,
            )
1816
1817

        if has_decode:
1818
1819
            assert attn_metadata.decode is not None
            decode_q_nope, decode_q_pe = decode_q.split(
1820
1821
                [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
            )
1822
1823
            # Convert from (B, N, P) to (N, B, P)
            decode_q_nope = decode_q_nope.transpose(0, 1)
1824

1825
1826
1827
            # Pads the head_dim if necessary (for the underlying kernel)
            if self.q_pad_num_heads is not None:
                B, N, L = decode_q_pe.shape
1828
                decode_pe_padded = decode_q_pe.new_empty((B, self.q_pad_num_heads, L))
1829
1830
1831
1832
                decode_pe_padded.resize_((B, N, L))
                decode_pe_padded.copy_(decode_q_pe)
                decode_q_pe = decode_pe_padded

1833
1834
            if is_rocm_aiter_fp8bmm_enabled():
                # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)
1835
1836
1837
1838
1839
1840
1841
                decode_ql_nope = aiter_triton_fp8_bmm(
                    decode_q_nope,
                    self.W_K,
                    self.W_K_scale,
                    group_size=128,
                    transpose_bm=True,
                )
1842
            else:
1843
1844
1845
1846
1847
                # Pads the head_dim if necessary (for the underlying kernel)
                N, B, P = decode_q_nope.shape
                _, _, L = self.W_UK_T.shape
                if self.q_pad_num_heads is not None:
                    decode_ql_nope = decode_q_nope.new_empty(
1848
1849
                        (self.q_pad_num_heads, B, L)
                    )
1850
1851
1852
1853
1854
                    decode_ql_nope.resize_((N, B, L))

                else:
                    decode_ql_nope = decode_q_nope.new_empty((N, B, L))

1855
                # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1856
                torch.bmm(decode_q_nope, self.W_UK_T, out=decode_ql_nope)
1857
1858
                # Convert from (N, B, L) to (B, N, L)
                decode_ql_nope = decode_ql_nope.transpose(0, 1)
1859

1860
1861
1862
            if fp8_attention:
                ql_nope_shape = decode_ql_nope.shape
                decode_ql_nope, _ = ops.scaled_fp8_quant(
1863
1864
1865
1866
1867
                    decode_ql_nope.reshape(
                        [ql_nope_shape[0], ql_nope_shape[1] * ql_nope_shape[2]]
                    ),
                    layer._q_scale,
                )
1868
1869
1870
                decode_ql_nope = decode_ql_nope.reshape(ql_nope_shape)
                q_pe_shape = decode_q_pe.shape
                decode_q_pe, _ = ops.scaled_fp8_quant(
1871
1872
1873
                    decode_q_pe.reshape([q_pe_shape[0], q_pe_shape[1] * q_pe_shape[2]]),
                    layer._q_scale,
                )
1874
1875
                decode_q_pe = decode_q_pe.reshape(q_pe_shape)

1876
1877
1878
1879
1880
1881
1882
1883
1884
            decode_q = (decode_ql_nope, decode_q_pe)
            if self.dcp_world_size > 1:
                assert not fp8_attention, "DCP not support fp8 kvcache now."
                # concatenate decode_ql_nope and decode_q_pe -> (B, N, L + P)
                decode_q = torch.cat(decode_q, dim=-1)
                # decode_q do allgather in head dim.
                decode_q = get_dcp_group().all_gather(decode_q, dim=1)

            # call decode attn
1885
1886
1887
            attn_out, lse = self._forward_decode(
                decode_q, kv_cache, attn_metadata, layer
            )
1888
1889
1890
1891
1892
1893

            # recorect dcp attn_out with lse.
            if self.dcp_world_size > 1:
                attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group())

            # v_up projection
1894
            self._v_up_proj(attn_out, out=output[:num_decode_tokens])
1895
        return output_padded