[Kimi-K3] Overlap low-M TP8 KDA projections (#54697)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Yongye Zhu
co-authored by
OpenAI Codex
OpenAI Codex
parent
12b9573c98
commit
3ba9907a1d
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Compare Kimi-K3 TP8 KDA sequential and multistream projections on B300."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
|
||||
shape_dynamic_skinny_gemm,
|
||||
)
|
||||
from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
|
||||
KDA_PROJECTION_OVERLAP_MAX_TOKENS,
|
||||
run_kda_projection_overlap,
|
||||
try_low_latency_gemm,
|
||||
)
|
||||
|
||||
ProjectionOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
|
||||
|
||||
def _capture(
|
||||
fn: Callable[[], list[ProjectionOutput]],
|
||||
) -> tuple[torch.cuda.CUDAGraph, list[ProjectionOutput]]:
|
||||
capture_stream = torch.cuda.Stream()
|
||||
capture_stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(capture_stream):
|
||||
fn()
|
||||
fn()
|
||||
capture_stream.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=capture_stream):
|
||||
outputs = fn()
|
||||
torch.cuda.current_stream().wait_stream(capture_stream)
|
||||
torch.accelerator.synchronize()
|
||||
return graph, outputs
|
||||
|
||||
|
||||
def _benchmark_graph(
|
||||
graph: torch.cuda.CUDAGraph,
|
||||
num_weight_sets: int,
|
||||
repeats: int,
|
||||
) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(repeats):
|
||||
graph.replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) * 1000 / repeats / num_weight_sets
|
||||
|
||||
|
||||
def _validate(
|
||||
output: ProjectionOutput,
|
||||
hidden_states: torch.Tensor,
|
||||
packed_weight: torch.Tensor,
|
||||
f_b_weight: torch.Tensor,
|
||||
) -> None:
|
||||
expected_qkvg = torch.nn.functional.linear(
|
||||
hidden_states.float(), packed_weight[:6144].float()
|
||||
)
|
||||
expected_fab = torch.nn.functional.linear(
|
||||
hidden_states.float(), packed_weight[6144:6284].float()
|
||||
)
|
||||
expected_g1 = torch.nn.functional.linear(
|
||||
expected_fab[:, :128].to(torch.bfloat16).float(),
|
||||
f_b_weight.float(),
|
||||
)
|
||||
for actual, expected in zip(
|
||||
output,
|
||||
(expected_qkvg, expected_g1, expected_fab[:, 128:]),
|
||||
):
|
||||
cosine = torch.nn.functional.cosine_similarity(
|
||||
actual.float().flatten(), expected.flatten(), dim=0
|
||||
).item()
|
||||
if cosine <= 0.999:
|
||||
raise AssertionError(f"projection cosine similarity is {cosine}")
|
||||
|
||||
|
||||
def _sequential_projection(
|
||||
hidden_states: torch.Tensor,
|
||||
packed_weight: torch.Tensor,
|
||||
f_b_weight: torch.Tensor,
|
||||
) -> ProjectionOutput:
|
||||
projected = try_low_latency_gemm(hidden_states, packed_weight)
|
||||
if projected is None:
|
||||
projected = torch.mm(hidden_states, packed_weight.t())
|
||||
f_a = projected[:, 6144:6272]
|
||||
g1 = torch.mm(f_a, f_b_weight.t())
|
||||
return projected[:, :6144], g1, projected[:, 6272:6284]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--tokens", type=int, nargs="+", default=range(1, 17))
|
||||
parser.add_argument("--weight-sets", type=int, default=4)
|
||||
parser.add_argument("--repeats", type=int, default=300)
|
||||
parser.add_argument("--trials", type=int, default=5)
|
||||
args = parser.parse_args()
|
||||
|
||||
if torch.cuda.get_device_capability() != (10, 3):
|
||||
raise SystemExit("This benchmark requires an SM103 GPU")
|
||||
if not shape_dynamic_skinny_gemm.is_available():
|
||||
raise SystemExit("This benchmark requires CuTe DSL")
|
||||
required_ops = ("dsv3_fused_a_gemm",)
|
||||
if not all(hasattr(torch.ops._C, name) for name in required_ops):
|
||||
raise SystemExit("This benchmark requires the KDA projection GEMMs")
|
||||
if not all(1 <= num_tokens <= 16 for num_tokens in args.tokens):
|
||||
raise SystemExit("--tokens must be in [1, 16]")
|
||||
|
||||
torch.manual_seed(42)
|
||||
packed_weights = [
|
||||
torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda")
|
||||
for _ in range(args.weight_sets)
|
||||
]
|
||||
f_b_weights = [
|
||||
torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
|
||||
for _ in range(args.weight_sets)
|
||||
]
|
||||
|
||||
from flashinfer.autotuner import autotune
|
||||
from flashinfer.gemm import mm_bf16
|
||||
|
||||
with torch.inference_mode(), autotune(tune_mode=True):
|
||||
mm_bf16(
|
||||
torch.empty(
|
||||
(KDA_PROJECTION_OVERLAP_MAX_TOKENS, 7168),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
),
|
||||
packed_weights[0][:6144].t(),
|
||||
pdl=True,
|
||||
backend="cute-dsl",
|
||||
)
|
||||
|
||||
print("M sequential (us) multistream (us) fastest")
|
||||
for num_tokens in args.tokens:
|
||||
hidden_states = [
|
||||
torch.randn(
|
||||
num_tokens,
|
||||
7168,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
for _ in range(args.weight_sets)
|
||||
]
|
||||
|
||||
def sequential(
|
||||
hidden_states: list[torch.Tensor] = hidden_states,
|
||||
) -> list[ProjectionOutput]:
|
||||
return [
|
||||
_sequential_projection(x, weight, f_b_weight)
|
||||
for x, weight, f_b_weight in zip(
|
||||
hidden_states, packed_weights, f_b_weights
|
||||
)
|
||||
]
|
||||
|
||||
candidates: list[tuple[str, Callable[[], list[ProjectionOutput]]]] = [
|
||||
("sequential", sequential)
|
||||
]
|
||||
if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
|
||||
aux_stream = torch.cuda.Stream()
|
||||
events = [
|
||||
(torch.cuda.Event(), torch.cuda.Event())
|
||||
for _ in range(args.weight_sets)
|
||||
]
|
||||
|
||||
def overlap(
|
||||
hidden_states: list[torch.Tensor] = hidden_states,
|
||||
aux_stream: torch.cuda.Stream = aux_stream,
|
||||
events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = events,
|
||||
) -> list[ProjectionOutput]:
|
||||
return [
|
||||
run_kda_projection_overlap(x, weight, f_b_weight, aux_stream, pair)
|
||||
for x, weight, f_b_weight, pair in zip(
|
||||
hidden_states, packed_weights, f_b_weights, events
|
||||
)
|
||||
]
|
||||
|
||||
candidates.append(("multistream", overlap))
|
||||
|
||||
graphs = {}
|
||||
for name, candidate in candidates:
|
||||
graph, outputs = _capture(candidate)
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
_validate(outputs[0], hidden_states[0], packed_weights[0], f_b_weights[0])
|
||||
graphs[name] = graph
|
||||
|
||||
samples = {name: [] for name in graphs}
|
||||
graph_items = list(graphs.items())
|
||||
for trial in range(args.trials):
|
||||
ordered_graphs = graph_items if trial % 2 == 0 else graph_items[::-1]
|
||||
for name, graph in ordered_graphs:
|
||||
samples[name].append(
|
||||
_benchmark_graph(graph, args.weight_sets, args.repeats)
|
||||
)
|
||||
results = {name: statistics.median(values) for name, values in samples.items()}
|
||||
|
||||
sequential_us = results["sequential"]
|
||||
multistream_us = results.get("multistream")
|
||||
multistream_text = (
|
||||
f"{multistream_us:18.3f}"
|
||||
if multistream_us is not None
|
||||
else " -"
|
||||
)
|
||||
fastest = (
|
||||
"multistream"
|
||||
if multistream_us is not None and multistream_us < sequential_us
|
||||
else "sequential"
|
||||
)
|
||||
print(f"{num_tokens:2d} {sequential_us:15.3f}{multistream_text} {fastest}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -253,6 +253,65 @@ def test_residual_cute_configs_match_measured_table() -> None:
|
||||
assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS
|
||||
|
||||
|
||||
def test_kda_overlap_configs_match_measured_table() -> None:
|
||||
assert k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS == 14
|
||||
assert k3_gemm.KDA_SKINNY_N_MAX_TOKENS == 14
|
||||
assert k3_gemm.KDA_SKINNY_K_MAX_TOKENS == 14
|
||||
assert _config_tuple(k3_gemm.KDA_M1_QKVG_CONFIG) == (64, 4, 2, 8)
|
||||
assert _config_tuple(k3_gemm.KDA_M1_FAB_CONFIG) == (224, 1, 2, 8)
|
||||
assert {
|
||||
num_tokens: _config_tuple(config)
|
||||
for num_tokens, config in k3_gemm.KDA_QKVG_CONFIGS.items()
|
||||
} == {
|
||||
1: (64, 4, 2, 8),
|
||||
2: (64, 3, 2, 8),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
|
||||
def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
|
||||
from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
|
||||
|
||||
kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
|
||||
nn.Module.__init__(kda)
|
||||
kda.tp_size = tp_size
|
||||
kda._projection_overlap_max_tokens = 0
|
||||
|
||||
assert not k3_gemm._enable_kda_projection_overlap(kda)
|
||||
assert kda._projection_overlap_max_tokens == 0
|
||||
|
||||
|
||||
def test_kda_qkvg_autotune_enables_full_overlap(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import flashinfer.gemm
|
||||
|
||||
from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
|
||||
|
||||
kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
|
||||
nn.Module.__init__(kda)
|
||||
kda._projection_overlap_max_tokens = max(k3_gemm.KDA_QKVG_CONFIGS)
|
||||
kda.in_proj_qkvgfab = nn.Module()
|
||||
kda.in_proj_qkvgfab.weight = nn.Parameter(
|
||||
torch.empty(6144, 8, dtype=torch.bfloat16),
|
||||
requires_grad=False,
|
||||
)
|
||||
calls = []
|
||||
|
||||
def fake_mm_bf16(a, b, *, pdl, backend):
|
||||
calls.append((a.shape, b.shape, pdl, backend))
|
||||
return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype)
|
||||
|
||||
monkeypatch.setattr(flashinfer.gemm, "mm_bf16", fake_mm_bf16)
|
||||
|
||||
k3_gemm.autotune_kda_qkvg(kda)
|
||||
|
||||
assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")]
|
||||
assert (
|
||||
kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS
|
||||
)
|
||||
|
||||
|
||||
def test_glm52_projection_plans_are_separate() -> None:
|
||||
qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION
|
||||
q_b = glm52_gemm.GLM52_Q_B_PROJECTION
|
||||
@@ -812,6 +871,65 @@ def test_cute_selected_shapes(
|
||||
assert cosine > 0.999
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14])
|
||||
def test_kda_projection_overlap_cuda_graph(num_tokens: int) -> None:
|
||||
"""The TP8 split must preserve projection results across both streams."""
|
||||
_require_sm103_and_cute()
|
||||
_require_sm103_and_dsv3()
|
||||
torch.manual_seed(43)
|
||||
hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
|
||||
packed_weight = torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda")
|
||||
f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
|
||||
aux_stream = torch.cuda.Stream()
|
||||
events = (torch.cuda.Event(), torch.cuda.Event())
|
||||
|
||||
k3_gemm.run_kda_projection_overlap(
|
||||
hidden_states,
|
||||
packed_weight,
|
||||
f_b_weight,
|
||||
aux_stream,
|
||||
events,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
capture_stream = torch.cuda.Stream()
|
||||
capture_stream.wait_stream(torch.cuda.current_stream())
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=capture_stream):
|
||||
projected_qkvg, g1, beta = k3_gemm.run_kda_projection_overlap(
|
||||
hidden_states,
|
||||
packed_weight,
|
||||
f_b_weight,
|
||||
aux_stream,
|
||||
events,
|
||||
)
|
||||
torch.cuda.current_stream().wait_stream(capture_stream)
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
assert beta.stride() == ((140 if num_tokens == 1 else 144), 1)
|
||||
|
||||
expected_qkvg = torch.nn.functional.linear(
|
||||
hidden_states.float(), packed_weight[:6144].float()
|
||||
)
|
||||
expected_fab = torch.nn.functional.linear(
|
||||
hidden_states.float(), packed_weight[6144:6284].float()
|
||||
)
|
||||
expected_g1 = torch.nn.functional.linear(
|
||||
expected_fab[:, :128].to(torch.bfloat16).float(),
|
||||
f_b_weight.float(),
|
||||
)
|
||||
for actual, expected in (
|
||||
(projected_qkvg, expected_qkvg),
|
||||
(g1, expected_g1),
|
||||
(beta, expected_fab[:, 128:]),
|
||||
):
|
||||
cosine = torch.nn.functional.cosine_similarity(
|
||||
actual.float().flatten(), expected.flatten(), dim=0
|
||||
).item()
|
||||
assert cosine > 0.999
|
||||
|
||||
|
||||
def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
|
||||
"""Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch."""
|
||||
if not tokens:
|
||||
|
||||
@@ -114,6 +114,12 @@ def _warmup_kimi_k3_gemm_rs_ar() -> None:
|
||||
logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled)
|
||||
|
||||
|
||||
def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None:
|
||||
module = sys.modules.get("vllm.models.kimi_k3.nvidia.low_latency_gemm")
|
||||
if module is not None:
|
||||
module.autotune_kda_qkvg(model)
|
||||
|
||||
|
||||
def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
|
||||
from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
|
||||
minimax_m3_msa_warmup,
|
||||
@@ -343,6 +349,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
|
||||
):
|
||||
_run_flashinfer_autotune_dummy_runs(runner)
|
||||
replayssm_autotune_warmup(runner)
|
||||
_autotune_kimi_k3_kda_qkvg(runner.get_model())
|
||||
finally:
|
||||
set_autotune_process_group(None)
|
||||
|
||||
|
||||
@@ -399,6 +399,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
config: KimiLinearConfig,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
aux_stream: torch.cuda.Stream | None = None,
|
||||
run_gemm_rs_ar: bool = False,
|
||||
) -> None:
|
||||
super().__init__(config, vllm_config, prefix)
|
||||
@@ -421,6 +422,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
assert kda_config.get("use_full_rank_gate", False), (
|
||||
"KimiK3DeltaAttention requires a full-rank gate"
|
||||
)
|
||||
self._projection_aux_stream = aux_stream
|
||||
self._projection_events = (
|
||||
(torch.cuda.Event(), torch.cuda.Event()) if aux_stream is not None else None
|
||||
)
|
||||
self._projection_overlap_max_tokens = 0
|
||||
|
||||
# Keep f_a before the narrow beta shard, then pad each TP-local row
|
||||
# to select the aligned BF16 GEMM path.
|
||||
@@ -605,19 +611,44 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
positions: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
num_tokens = hidden_states.size(0)
|
||||
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
|
||||
split_sizes = [
|
||||
3 * self.local_projection_size,
|
||||
self.local_projection_size,
|
||||
self.head_dim,
|
||||
self.local_num_heads,
|
||||
]
|
||||
if self.in_proj_padding:
|
||||
split_sizes.append(self.in_proj_padding)
|
||||
projected = projected_qkvgfab.split(split_sizes, dim=-1)
|
||||
mixed_qkv, g_proj_states, f_a, beta = projected[:4]
|
||||
projection_events = self._projection_events
|
||||
projection_aux_stream = self._projection_aux_stream
|
||||
if (
|
||||
0 < num_tokens <= self._projection_overlap_max_tokens
|
||||
and hidden_states.stride() == (self.hidden_size, 1)
|
||||
and projection_events is not None
|
||||
and projection_aux_stream is not None
|
||||
and torch.cuda.is_current_stream_capturing()
|
||||
):
|
||||
from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
|
||||
run_kda_projection_overlap,
|
||||
)
|
||||
|
||||
projected_qkvg, g1, beta = run_kda_projection_overlap(
|
||||
hidden_states,
|
||||
self.in_proj_qkvgfab.weight,
|
||||
self.f_b_proj.weight,
|
||||
projection_aux_stream,
|
||||
projection_events,
|
||||
)
|
||||
mixed_qkv, g_proj_states = projected_qkvg.split(
|
||||
[3 * self.local_projection_size, self.local_projection_size],
|
||||
dim=-1,
|
||||
)
|
||||
else:
|
||||
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
|
||||
split_sizes = [
|
||||
3 * self.local_projection_size,
|
||||
self.local_projection_size,
|
||||
self.head_dim,
|
||||
self.local_num_heads,
|
||||
]
|
||||
if self.in_proj_padding:
|
||||
split_sizes.append(self.in_proj_padding)
|
||||
projected = projected_qkvgfab.split(split_sizes, dim=-1)
|
||||
mixed_qkv, g_proj_states, f_a, beta = projected[:4]
|
||||
g1 = self.f_b_proj(f_a)[0]
|
||||
|
||||
g1 = self.f_b_proj(f_a)[0]
|
||||
beta = beta.unsqueeze(0)
|
||||
g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
|
||||
g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
|
||||
|
||||
@@ -35,12 +35,32 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
UnquantizedEmbeddingMethod,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
|
||||
|
||||
Backend = Literal["cute", "dsv3_fused_a"]
|
||||
# A resolved per-token-count call: the backend plus its CuTe config (None for
|
||||
# dsv3, which needs no config).
|
||||
ResolvedCall = tuple[Backend, SkinnyGemmConfig | None]
|
||||
|
||||
# TP8 KDA projection split, measured together under CUDA graph capture on B300.
|
||||
# Q/K/V/G stay on the graph's main stream while F_A/beta and F_B run on the
|
||||
# model's auxiliary stream.
|
||||
KDA_M1_QKVG_CONFIG = SkinnyGemmConfig(1, 64, 4, 2, 8)
|
||||
KDA_M1_FAB_CONFIG = SkinnyGemmConfig(1, 224, 1, 2, 8)
|
||||
KDA_QKVG_CONFIGS = {
|
||||
1: KDA_M1_QKVG_CONFIG,
|
||||
2: SkinnyGemmConfig(2, 64, 3, 2, 8),
|
||||
}
|
||||
# The captured end-to-end projection sweep wins through M=14 and regresses at
|
||||
# M=15 and M=16, where the original packed projection remains selected.
|
||||
KDA_PROJECTION_OVERLAP_MAX_TOKENS = 14
|
||||
KDA_SKINNY_N_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS
|
||||
KDA_SKINNY_K_MAX_TOKENS = KDA_PROJECTION_OVERLAP_MAX_TOKENS
|
||||
_KDA_QKVG_SIZE = 4 * 1536
|
||||
_KDA_FAB_SIZE = 128 + 12
|
||||
_KDA_PACKED_SIZE = 6288
|
||||
_KDA_TP_SIZE = 8
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProjectionSpec:
|
||||
@@ -683,6 +703,137 @@ def _run_plan(
|
||||
return output
|
||||
|
||||
|
||||
def run_kda_projection_overlap(
|
||||
hidden_states: torch.Tensor,
|
||||
packed_weight: torch.Tensor,
|
||||
f_b_weight: torch.Tensor,
|
||||
aux_stream: torch.cuda.Stream,
|
||||
events: tuple[torch.cuda.Event, torch.cuda.Event],
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Run the TP8 KDA decode projection branches concurrently.
|
||||
|
||||
Args:
|
||||
hidden_states: Packed BF16 input with shape ``[M, 7168]``.
|
||||
packed_weight: Existing Q/K/V/G/F_A/beta/pad weight with shape
|
||||
``[6288, 7168]``.
|
||||
f_b_weight: F_B weight with shape ``[1536, 128]``.
|
||||
aux_stream: Stream for the F_A/beta then F_B branch.
|
||||
events: Start and completion events for the stream fork and join.
|
||||
|
||||
Returns:
|
||||
The Q/K/V/G projection, F_B output, and beta projection.
|
||||
"""
|
||||
num_tokens = hidden_states.shape[0]
|
||||
qkvg_weight = packed_weight[:_KDA_QKVG_SIZE]
|
||||
|
||||
def run_qkvg() -> torch.Tensor:
|
||||
if config := KDA_QKVG_CONFIGS.get(num_tokens):
|
||||
return shape_dynamic_skinny_gemm(
|
||||
hidden_states,
|
||||
qkvg_weight,
|
||||
config,
|
||||
None,
|
||||
)
|
||||
if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
|
||||
from flashinfer.gemm import mm_bf16
|
||||
|
||||
return mm_bf16(
|
||||
hidden_states,
|
||||
qkvg_weight.t(),
|
||||
pdl=True,
|
||||
backend="cute-dsl",
|
||||
)
|
||||
return torch.mm(hidden_states, qkvg_weight.t())
|
||||
|
||||
def run_fab_fb() -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if num_tokens == 1:
|
||||
projected_fab = shape_dynamic_skinny_gemm(
|
||||
hidden_states,
|
||||
packed_weight[_KDA_QKVG_SIZE : _KDA_QKVG_SIZE + _KDA_FAB_SIZE],
|
||||
KDA_M1_FAB_CONFIG,
|
||||
None,
|
||||
)
|
||||
f_a, beta = projected_fab.split([128, 12], dim=-1)
|
||||
f_a = f_a.as_strided((1, 128), (128, 1))
|
||||
g1 = torch.empty(
|
||||
(1, 1536), dtype=hidden_states.dtype, device=hidden_states.device
|
||||
)
|
||||
ops.dsv3_fused_a_gemm(g1, f_a, f_b_weight.t(), enable_pdl=True)
|
||||
else:
|
||||
from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
|
||||
kda_skinny_gemm,
|
||||
)
|
||||
|
||||
if num_tokens <= KDA_SKINNY_N_MAX_TOKENS:
|
||||
projected_fab = kda_skinny_gemm.run_n(
|
||||
hidden_states,
|
||||
packed_weight[_KDA_QKVG_SIZE:],
|
||||
)
|
||||
else:
|
||||
projected_fab = torch.mm(
|
||||
hidden_states,
|
||||
packed_weight[_KDA_QKVG_SIZE:].t(),
|
||||
)
|
||||
beta = projected_fab[:, 128:140]
|
||||
if num_tokens <= KDA_SKINNY_K_MAX_TOKENS:
|
||||
g1 = kda_skinny_gemm.run_k(projected_fab, f_b_weight)
|
||||
else:
|
||||
g1 = torch.mm(projected_fab[:, :128], f_b_weight.t())
|
||||
return g1, beta
|
||||
|
||||
projected_qkvg, (g1, beta) = maybe_execute_in_parallel(
|
||||
run_qkvg,
|
||||
run_fab_fb,
|
||||
events[0],
|
||||
events[1],
|
||||
aux_stream,
|
||||
)
|
||||
return projected_qkvg, g1, beta
|
||||
|
||||
|
||||
def autotune_kda_qkvg(model: nn.Module) -> None:
|
||||
"""Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
|
||||
from flashinfer.gemm import mm_bf16
|
||||
|
||||
from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
|
||||
|
||||
children: list[KimiK3DeltaAttention] = []
|
||||
weights_by_shape: dict[
|
||||
tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
|
||||
] = {}
|
||||
for child in model.modules():
|
||||
if not isinstance(child, KimiK3DeltaAttention):
|
||||
continue
|
||||
if child._projection_overlap_max_tokens <= 0:
|
||||
continue
|
||||
children.append(child)
|
||||
qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
|
||||
shape = (
|
||||
KDA_PROJECTION_OVERLAP_MAX_TOKENS,
|
||||
qkvg_weight.shape[0],
|
||||
qkvg_weight.shape[1],
|
||||
qkvg_weight.dtype,
|
||||
qkvg_weight.device,
|
||||
)
|
||||
weights_by_shape.setdefault(shape, qkvg_weight)
|
||||
|
||||
for shape, qkvg_weight in weights_by_shape.items():
|
||||
num_tokens = shape[0]
|
||||
hidden_states = torch.empty(
|
||||
(num_tokens, qkvg_weight.shape[1]),
|
||||
dtype=qkvg_weight.dtype,
|
||||
device=qkvg_weight.device,
|
||||
)
|
||||
mm_bf16(
|
||||
hidden_states,
|
||||
qkvg_weight.t(),
|
||||
pdl=True,
|
||||
backend="cute-dsl",
|
||||
)
|
||||
for child in children:
|
||||
child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS
|
||||
|
||||
|
||||
def _run_residual_plan(
|
||||
residual_plan: dict[int, SkinnyGemmConfig],
|
||||
x: torch.Tensor,
|
||||
@@ -777,6 +928,40 @@ class KimiK3LowLatencyEmbeddingMethod(
|
||||
pass
|
||||
|
||||
|
||||
def _enable_kda_projection_overlap(module: nn.Module) -> bool:
|
||||
from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
|
||||
|
||||
if envs.VLLM_BATCH_INVARIANT:
|
||||
return False
|
||||
|
||||
enabled = False
|
||||
for child in module.modules():
|
||||
if not isinstance(child, KimiK3DeltaAttention):
|
||||
continue
|
||||
if child.tp_size != _KDA_TP_SIZE:
|
||||
continue
|
||||
packed_weight = child.in_proj_qkvgfab.weight
|
||||
f_b_weight = child.f_b_proj.weight
|
||||
if (
|
||||
type(child.in_proj_qkvgfab.quant_method) is not KimiK3LowLatencyLinearMethod
|
||||
or type(child.f_b_proj.quant_method) is not KimiK3LowLatencyLinearMethod
|
||||
or child._projection_aux_stream is None
|
||||
or child._projection_events is None
|
||||
or packed_weight.shape != (_KDA_PACKED_SIZE, 7168)
|
||||
or f_b_weight.shape != (1536, 128)
|
||||
or packed_weight.dtype != torch.bfloat16
|
||||
or f_b_weight.dtype != torch.bfloat16
|
||||
or not packed_weight.is_cuda
|
||||
or not f_b_weight.is_cuda
|
||||
or not packed_weight.is_contiguous()
|
||||
or not f_b_weight.is_contiguous()
|
||||
):
|
||||
continue
|
||||
child._projection_overlap_max_tokens = max(KDA_QKVG_CONFIGS)
|
||||
enabled = True
|
||||
return enabled
|
||||
|
||||
|
||||
def enable_kimi_k3_low_latency_gemm(
|
||||
module: nn.Module,
|
||||
dtype: torch.dtype,
|
||||
@@ -825,6 +1010,16 @@ def enable_kimi_k3_low_latency_gemm(
|
||||
# TP8 deployment does not compile TP4 configs and vice versa.
|
||||
warmup_configs.update(config for _, config in spec.cute_configs)
|
||||
residual_warmup_configs.update(config for _, config in spec.residual_configs)
|
||||
if _enable_kda_projection_overlap(module):
|
||||
from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
|
||||
kda_skinny_gemm,
|
||||
)
|
||||
|
||||
warmup_configs.update((*KDA_QKVG_CONFIGS.values(), KDA_M1_FAB_CONFIG))
|
||||
kda_skinny_gemm.request_warmup(
|
||||
set(range(2, KDA_SKINNY_N_MAX_TOKENS + 1)),
|
||||
set(range(2, KDA_SKINNY_K_MAX_TOKENS + 1)),
|
||||
)
|
||||
|
||||
if shape_dynamic_skinny_gemm.is_available():
|
||||
if warmup_configs:
|
||||
|
||||
@@ -895,6 +895,7 @@ class KimiDecoderLayer(nn.Module):
|
||||
config,
|
||||
vllm_config,
|
||||
prefix=f"{prefix}.self_attn",
|
||||
aux_stream=aux_stream,
|
||||
run_gemm_rs_ar=run_gemm_rs_ar,
|
||||
)
|
||||
self._self_attn_writes_output = False
|
||||
|
||||
@@ -0,0 +1,507 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from
|
||||
# https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh
|
||||
"""Kimi-K3 TP8 skinny GEMMs for the KDA F_A/beta and F_B projections."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import torch
|
||||
from cutlass import BFloat16, Float32
|
||||
from cutlass._mlir.dialects import llvm
|
||||
from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
|
||||
from cutlass.cutlass_dsl import T, dsl_user_op
|
||||
|
||||
_FAB_N = 144
|
||||
_FAB_K = 7168
|
||||
_FAB_VECTOR_WIDTH = 16
|
||||
_FAB_BLOCK_SIZE = _FAB_K // _FAB_VECTOR_WIDTH
|
||||
_FAB_NUM_WARPS = _FAB_BLOCK_SIZE // cute.arch.WARP_SIZE
|
||||
|
||||
_FB_N = 1536
|
||||
_FB_K = 128
|
||||
_FB_VECTOR_WIDTH = 8
|
||||
_FB_K_LANES = _FB_K // _FB_VECTOR_WIDTH
|
||||
|
||||
SKINNY_N_SPLIT_N = {
|
||||
2: 1,
|
||||
3: 2,
|
||||
4: 3,
|
||||
5: 3,
|
||||
6: 3,
|
||||
7: 3,
|
||||
8: 3,
|
||||
9: 2,
|
||||
10: 2,
|
||||
11: 2,
|
||||
12: 2,
|
||||
13: 2,
|
||||
14: 1,
|
||||
15: 1,
|
||||
16: 1,
|
||||
}
|
||||
|
||||
SKINNY_K_SPLIT_N = {
|
||||
2: 8,
|
||||
3: 16,
|
||||
4: 16,
|
||||
5: 24,
|
||||
6: 24,
|
||||
7: 24,
|
||||
8: 24,
|
||||
9: 12,
|
||||
10: 16,
|
||||
11: 16,
|
||||
12: 12,
|
||||
13: 8,
|
||||
14: 8,
|
||||
15: 8,
|
||||
16: 8,
|
||||
}
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def _fma_f32_bf16(
|
||||
a: BFloat16,
|
||||
b: BFloat16,
|
||||
acc: Float32,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Float32:
|
||||
a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
|
||||
b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
|
||||
result = llvm.inline_asm(
|
||||
T.f32(),
|
||||
[a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
|
||||
"fma.rn.f32.bf16 $0, $1, $2, $3;",
|
||||
"=f,h,h,f",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return Float32(result)
|
||||
|
||||
|
||||
class _KdaSkinnyNGemm:
|
||||
def __init__(self, num_rows: int, split_n: int) -> None:
|
||||
if _FAB_N % split_n:
|
||||
raise ValueError("split_n must divide the FAB output size")
|
||||
if num_rows * split_n > _FAB_BLOCK_SIZE:
|
||||
raise ValueError("the FAB output tile must fit in one block")
|
||||
self.num_rows = num_rows
|
||||
self.split_n = split_n
|
||||
|
||||
@cute.jit
|
||||
def __call__(
|
||||
self,
|
||||
a: cute.Tensor,
|
||||
weight: cute.Tensor,
|
||||
output: cute.Tensor,
|
||||
stream: cuda.CUstream,
|
||||
) -> None:
|
||||
copy_a = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyG2ROp(),
|
||||
BFloat16,
|
||||
num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width,
|
||||
load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
|
||||
)
|
||||
copy_weight = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyG2ROp(),
|
||||
BFloat16,
|
||||
num_bits_per_copy=_FAB_VECTOR_WIDTH * BFloat16.width,
|
||||
load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
|
||||
)
|
||||
self.kernel(a, weight, output, copy_a, copy_weight).launch(
|
||||
grid=[_FAB_N // self.split_n, 1, 1],
|
||||
block=[_FAB_BLOCK_SIZE, 1, 1],
|
||||
smem=_FAB_NUM_WARPS * self.num_rows * self.split_n * 4,
|
||||
stream=stream,
|
||||
use_pdl=True,
|
||||
min_blocks_per_mp=1,
|
||||
)
|
||||
|
||||
@cute.kernel
|
||||
def kernel(
|
||||
self,
|
||||
a: cute.Tensor,
|
||||
weight: cute.Tensor,
|
||||
output: cute.Tensor,
|
||||
copy_a: cute.CopyAtom,
|
||||
copy_weight: cute.CopyAtom,
|
||||
) -> None:
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
block_idx, _, _ = cute.arch.block_idx()
|
||||
warp_idx = cute.arch.warp_idx()
|
||||
num_rows: cutlass.Constexpr = self.num_rows
|
||||
split_n: cutlass.Constexpr = self.split_n
|
||||
|
||||
a_vectors = cute.logical_divide(a, (None, _FAB_VECTOR_WIDTH))
|
||||
weight_vectors = cute.logical_divide(weight, (None, _FAB_VECTOR_WIDTH))
|
||||
a_tiles = cute.logical_divide(
|
||||
a_vectors,
|
||||
(None, (None, _FAB_BLOCK_SIZE)),
|
||||
)
|
||||
weight_tiles = cute.logical_divide(
|
||||
weight_vectors,
|
||||
(None, (None, _FAB_BLOCK_SIZE)),
|
||||
)
|
||||
a_tile = a_tiles[None, (None, (tidx, None))]
|
||||
n_base = block_idx * split_n
|
||||
|
||||
a_regs = cute.make_rmem_tensor(
|
||||
cute.make_layout(
|
||||
(num_rows, _FAB_VECTOR_WIDTH),
|
||||
stride=(_FAB_VECTOR_WIDTH, 1),
|
||||
),
|
||||
BFloat16,
|
||||
)
|
||||
weight_regs = cute.make_rmem_tensor(
|
||||
cute.make_layout(
|
||||
(split_n, _FAB_VECTOR_WIDTH),
|
||||
stride=(_FAB_VECTOR_WIDTH, 1),
|
||||
),
|
||||
BFloat16,
|
||||
)
|
||||
acc = cute.make_rmem_tensor(
|
||||
cute.make_layout((num_rows, split_n), stride=(split_n, 1)),
|
||||
Float32,
|
||||
)
|
||||
acc.fill(0.0)
|
||||
|
||||
for ni in cutlass.range_constexpr(split_n):
|
||||
weight_tile = weight_tiles[
|
||||
n_base + ni,
|
||||
(None, (tidx, None)),
|
||||
]
|
||||
cute.copy(copy_weight, weight_tile[None, 0], weight_regs[ni, None])
|
||||
|
||||
cute.arch.griddepcontrol_wait()
|
||||
|
||||
for mi in cutlass.range_constexpr(num_rows):
|
||||
cute.copy(copy_a, a_tile[mi, None, 0], a_regs[mi, None])
|
||||
for mi in cutlass.range_constexpr(num_rows):
|
||||
for ni in cutlass.range_constexpr(split_n):
|
||||
for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH):
|
||||
acc[mi, ni] = _fma_f32_bf16(
|
||||
a_regs[mi, vi],
|
||||
weight_regs[ni, vi],
|
||||
acc[mi, ni],
|
||||
)
|
||||
acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
|
||||
|
||||
partials_layout = cute.make_layout(
|
||||
(_FAB_NUM_WARPS, num_rows, split_n),
|
||||
stride=(num_rows * split_n, split_n, 1),
|
||||
)
|
||||
smem = cutlass.utils.SmemAllocator()
|
||||
partials = smem.allocate_tensor(
|
||||
Float32,
|
||||
partials_layout,
|
||||
byte_alignment=16,
|
||||
)
|
||||
with cute.arch.elect_one():
|
||||
for mi in cutlass.range_constexpr(num_rows):
|
||||
for ni in cutlass.range_constexpr(split_n):
|
||||
partials[warp_idx, mi, ni] = acc[mi, ni]
|
||||
|
||||
cute.arch.griddepcontrol_launch_dependents()
|
||||
cute.arch.sync_threads()
|
||||
|
||||
if tidx < num_rows * split_n:
|
||||
mi = tidx // split_n
|
||||
ni = tidx % split_n
|
||||
total = (
|
||||
partials[None, mi, ni]
|
||||
.load()
|
||||
.reduce(
|
||||
cute.ReductionOp.ADD,
|
||||
init_val=Float32(0.0),
|
||||
reduction_profile=0,
|
||||
)
|
||||
)
|
||||
output[mi, n_base + ni] = Float32(total).to(BFloat16)
|
||||
|
||||
|
||||
class _KdaSkinnyKGemm:
|
||||
def __init__(self, num_rows: int, split_n: int) -> None:
|
||||
if _FB_N % split_n:
|
||||
raise ValueError("split_n must divide the FB output size")
|
||||
if split_n * _FB_K_LANES > 1024:
|
||||
raise ValueError("the FB output tile exceeds the block-size limit")
|
||||
self.num_rows = num_rows
|
||||
self.split_n = split_n
|
||||
|
||||
@cute.jit
|
||||
def __call__(
|
||||
self,
|
||||
a: cute.Tensor,
|
||||
weight: cute.Tensor,
|
||||
output: cute.Tensor,
|
||||
stream: cuda.CUstream,
|
||||
) -> None:
|
||||
copy_atom = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyG2ROp(),
|
||||
BFloat16,
|
||||
num_bits_per_copy=_FB_VECTOR_WIDTH * BFloat16.width,
|
||||
load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
|
||||
)
|
||||
self.kernel(a, weight, output, copy_atom).launch(
|
||||
grid=[_FB_N // self.split_n, 1, 1],
|
||||
block=[self.split_n * _FB_K_LANES, 1, 1],
|
||||
stream=stream,
|
||||
use_pdl=True,
|
||||
min_blocks_per_mp=1,
|
||||
)
|
||||
|
||||
@cute.kernel
|
||||
def kernel(
|
||||
self,
|
||||
a: cute.Tensor,
|
||||
weight: cute.Tensor,
|
||||
output: cute.Tensor,
|
||||
copy_atom: cute.CopyAtom,
|
||||
) -> None:
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
block_idx, _, _ = cute.arch.block_idx()
|
||||
num_rows: cutlass.Constexpr = self.num_rows
|
||||
split_n: cutlass.Constexpr = self.split_n
|
||||
k_lane = tidx % _FB_K_LANES
|
||||
n_idx = block_idx * split_n + tidx // _FB_K_LANES
|
||||
|
||||
a_vectors = cute.logical_divide(a, (None, _FB_VECTOR_WIDTH))
|
||||
weight_vectors = cute.logical_divide(weight, (None, _FB_VECTOR_WIDTH))
|
||||
a_tiles = cute.logical_divide(
|
||||
a_vectors,
|
||||
(None, (None, _FB_K_LANES)),
|
||||
)
|
||||
weight_tiles = cute.logical_divide(
|
||||
weight_vectors,
|
||||
(None, (None, _FB_K_LANES)),
|
||||
)
|
||||
a_tile = a_tiles[None, (None, (k_lane, None))]
|
||||
weight_tile = weight_tiles[n_idx, (None, (k_lane, None))]
|
||||
weight_regs = cute.make_rmem_tensor(_FB_VECTOR_WIDTH, BFloat16)
|
||||
input_regs = cute.make_rmem_tensor(
|
||||
cute.make_layout(
|
||||
(num_rows, _FB_VECTOR_WIDTH),
|
||||
stride=(_FB_VECTOR_WIDTH, 1),
|
||||
),
|
||||
BFloat16,
|
||||
)
|
||||
cute.copy(
|
||||
copy_atom,
|
||||
weight_tile[None, 0],
|
||||
weight_regs,
|
||||
)
|
||||
|
||||
cute.arch.griddepcontrol_wait()
|
||||
|
||||
for mi in cutlass.range_constexpr(num_rows):
|
||||
cute.copy(
|
||||
copy_atom,
|
||||
a_tile[mi, None, 0],
|
||||
input_regs[mi, None],
|
||||
)
|
||||
for mi in cutlass.range_constexpr(num_rows):
|
||||
acc = Float32(0.0)
|
||||
for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH):
|
||||
acc = _fma_f32_bf16(
|
||||
input_regs[mi, vi],
|
||||
weight_regs[vi],
|
||||
acc,
|
||||
)
|
||||
for offset in (8, 4, 2, 1):
|
||||
acc += cute.arch.shuffle_sync_bfly(acc, offset)
|
||||
output[mi, n_idx] = acc.to(BFloat16)
|
||||
|
||||
cute.arch.griddepcontrol_launch_dependents()
|
||||
|
||||
|
||||
class KdaSkinnyGemm:
|
||||
def __init__(self) -> None:
|
||||
self._compiled_n: dict[tuple[int, int], Callable[..., None]] = {}
|
||||
self._compiled_k: dict[tuple[int, int], Callable[..., None]] = {}
|
||||
self._warmup_n: set[int] = set()
|
||||
self._warmup_k: set[int] = set()
|
||||
self._warmup_registered = False
|
||||
|
||||
def _compile_n(self, num_rows: int) -> None:
|
||||
device = torch.accelerator.current_device_index()
|
||||
key = (device, num_rows)
|
||||
if key in self._compiled_n:
|
||||
return
|
||||
split_n = SKINNY_N_SPLIT_N[num_rows]
|
||||
self._compiled_n[key] = cute.compile(
|
||||
_KdaSkinnyNGemm(num_rows, split_n),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(num_rows, _FAB_K),
|
||||
stride=(_FAB_K, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(_FAB_N, _FAB_K),
|
||||
stride=(_FAB_K, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(num_rows, _FAB_N),
|
||||
stride=(_FAB_N, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_stream(use_tvm_ffi_env_stream=True),
|
||||
options="--enable-tvm-ffi --ptxas-options -maxrregcount=255",
|
||||
)
|
||||
|
||||
def _compile_k(self, num_rows: int) -> None:
|
||||
device = torch.accelerator.current_device_index()
|
||||
key = (device, num_rows)
|
||||
if key in self._compiled_k:
|
||||
return
|
||||
split_n = SKINNY_K_SPLIT_N[num_rows]
|
||||
self._compiled_k[key] = cute.compile(
|
||||
_KdaSkinnyKGemm(num_rows, split_n),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(num_rows, _FB_K),
|
||||
stride=(_FAB_N, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(_FB_N, _FB_K),
|
||||
stride=(_FB_K, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_tensor(
|
||||
BFloat16,
|
||||
(num_rows, _FB_N),
|
||||
stride=(_FB_N, 1),
|
||||
assumed_align=32,
|
||||
),
|
||||
make_fake_stream(use_tvm_ffi_env_stream=True),
|
||||
options="--enable-tvm-ffi --ptxas-options -maxrregcount=255",
|
||||
)
|
||||
|
||||
def request_warmup(self, n_tokens: set[int], k_tokens: set[int]) -> None:
|
||||
if not n_tokens <= SKINNY_N_SPLIT_N.keys():
|
||||
raise ValueError("unsupported skinny-N warmup token count")
|
||||
if not k_tokens <= SKINNY_K_SPLIT_N.keys():
|
||||
raise ValueError("unsupported skinny-K warmup token count")
|
||||
self._warmup_n.update(n_tokens)
|
||||
self._warmup_k.update(k_tokens)
|
||||
if self._warmup_registered:
|
||||
return
|
||||
from vllm.model_executor.warmup.cutedsl_warmup import (
|
||||
register_cutedsl_warmup_provider,
|
||||
)
|
||||
|
||||
register_cutedsl_warmup_provider(self)
|
||||
self._warmup_registered = True
|
||||
|
||||
def get_cutedsl_warmup_compile_units(self):
|
||||
from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit
|
||||
|
||||
n_units = tuple(
|
||||
CuTeDSLCompileUnit(
|
||||
name="Kimi-K3 KDA skinny-N GEMM",
|
||||
key=("kimi-k3-kda-skinny-n", num_rows),
|
||||
compile=partial(self._compile_n, num_rows),
|
||||
)
|
||||
for num_rows in sorted(self._warmup_n)
|
||||
)
|
||||
k_units = tuple(
|
||||
CuTeDSLCompileUnit(
|
||||
name="Kimi-K3 KDA skinny-K GEMM",
|
||||
key=("kimi-k3-kda-skinny-k", num_rows),
|
||||
compile=partial(self._compile_k, num_rows),
|
||||
)
|
||||
for num_rows in sorted(self._warmup_k)
|
||||
)
|
||||
return n_units + k_units
|
||||
|
||||
@staticmethod
|
||||
def _validate(
|
||||
a: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
*,
|
||||
a_shape: tuple[int, int],
|
||||
weight_shape: tuple[int, int],
|
||||
) -> None:
|
||||
if a.shape != a_shape or weight.shape != weight_shape:
|
||||
raise ValueError(
|
||||
f"expected input {a_shape} and weight {weight_shape}, got "
|
||||
f"{tuple(a.shape)} and {tuple(weight.shape)}"
|
||||
)
|
||||
if a.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
|
||||
raise ValueError("KDA skinny GEMMs require BF16 input and weight")
|
||||
if not a.is_cuda or not weight.is_cuda or a.device != weight.device:
|
||||
raise ValueError("KDA skinny GEMM operands must share a CUDA device")
|
||||
if not a.is_contiguous() or not weight.is_contiguous():
|
||||
raise ValueError("KDA skinny GEMM operands must be contiguous")
|
||||
if a.data_ptr() % 32 or weight.data_ptr() % 32:
|
||||
raise ValueError("KDA skinny GEMM operands must be 32-byte aligned")
|
||||
|
||||
def run_n(self, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
num_rows = a.shape[0]
|
||||
if num_rows not in SKINNY_N_SPLIT_N:
|
||||
raise ValueError("skinny-N requires M in [2, 16]")
|
||||
self._validate(
|
||||
a,
|
||||
weight,
|
||||
a_shape=(num_rows, _FAB_K),
|
||||
weight_shape=(_FAB_N, _FAB_K),
|
||||
)
|
||||
output = torch.empty(
|
||||
(num_rows, _FAB_N),
|
||||
dtype=a.dtype,
|
||||
device=a.device,
|
||||
)
|
||||
with torch.accelerator.device_index(a.device.index):
|
||||
self._compile_n(num_rows)
|
||||
key = (a.device.index, num_rows)
|
||||
self._compiled_n[key](
|
||||
a,
|
||||
weight,
|
||||
output,
|
||||
)
|
||||
return output
|
||||
|
||||
def run_k(self, packed_fab: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
num_rows = packed_fab.shape[0]
|
||||
if num_rows not in SKINNY_K_SPLIT_N:
|
||||
raise ValueError("skinny-K requires M in [2, 16]")
|
||||
self._validate(
|
||||
packed_fab,
|
||||
weight,
|
||||
a_shape=(num_rows, _FAB_N),
|
||||
weight_shape=(_FB_N, _FB_K),
|
||||
)
|
||||
output = torch.empty(
|
||||
(num_rows, _FB_N),
|
||||
dtype=packed_fab.dtype,
|
||||
device=packed_fab.device,
|
||||
)
|
||||
with torch.accelerator.device_index(packed_fab.device.index):
|
||||
self._compile_k(num_rows)
|
||||
key = (packed_fab.device.index, num_rows)
|
||||
self._compiled_k[key](
|
||||
packed_fab[:, :_FB_K],
|
||||
weight,
|
||||
output,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
kda_skinny_gemm = KdaSkinnyGemm()
|
||||
Reference in New Issue
Block a user