[Perf][Qwen4Exp] Merge QSA QKVG and indexer QK projections (#59533)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: namgyu-youn <namgyu.dev@gmail.com>
This commit is contained in:
Thien Tran
co-authored by
Codex
namgyu-youn
Codex
namgyu-youn
parent
92044241a0
commit
042ab0305c
@@ -713,8 +713,8 @@ def test_low_latency_table_capability_routing(
|
||||
def test_qwen4_exp_hopper_plans_are_valid() -> None:
|
||||
plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS
|
||||
|
||||
assert len(plans) == 9
|
||||
assert sum(map(len, plans.values())) == 31
|
||||
assert len(plans) == 7
|
||||
assert sum(map(len, plans.values())) == 24
|
||||
assert (320, 10240) in plans
|
||||
assert (10240, 320) not in plans
|
||||
for (n, k), shape_plans in plans.items():
|
||||
@@ -728,10 +728,11 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"capability,expected_plans",
|
||||
[
|
||||
((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
|
||||
((10, 3), qwen4_exp_gemm.QWEN4_EXP_SM103_GEMM_PLANS),
|
||||
((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS),
|
||||
((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
|
||||
((12, 1), qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS),
|
||||
((12, 0), {}),
|
||||
((8, 0), {}),
|
||||
],
|
||||
)
|
||||
@@ -742,11 +743,16 @@ def test_qwen4_exp_gemm_capability_routing(
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
qwen4_exp_gemm.current_platform,
|
||||
"is_device_capability",
|
||||
lambda target: capability == target,
|
||||
"get_device_capability",
|
||||
lambda: capability,
|
||||
)
|
||||
|
||||
assert qwen4_exp_gemm._gemm_plans() == expected_plans
|
||||
assert (
|
||||
qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
qwen4_exp_gemm.current_platform.get_device_capability(), {}
|
||||
)
|
||||
== expected_plans
|
||||
)
|
||||
|
||||
|
||||
def test_installation_is_shape_specific_and_unquantized(
|
||||
|
||||
@@ -74,13 +74,14 @@ def _make_main_inputs(num_tokens: int, fp8_cache: bool) -> dict:
|
||||
slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens]
|
||||
slots[-1] = -1
|
||||
norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2
|
||||
qkv_width = 2 * (MAIN_HQ + MAIN_HK) * MAIN_D
|
||||
return dict(
|
||||
main_qkv=torch.randn(
|
||||
num_tokens,
|
||||
2 * (MAIN_HQ + MAIN_HK) * MAIN_D,
|
||||
qkv_width + (HQ + 1) * D,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
),
|
||||
)[:, :qkv_width],
|
||||
main_q_norm_weight=norm_weights[0],
|
||||
main_k_norm_weight=norm_weights[1],
|
||||
main_eps=EPS,
|
||||
@@ -316,9 +317,12 @@ def test_qsa_fused_prepare_matches_unfused(
|
||||
)
|
||||
unfused_compressed = fused_compressed.clone()
|
||||
|
||||
# TP=4 QKVG width makes the indexer slice row-strided, as in the merged
|
||||
# projection. The exact width is unimportant; the larger stride is tested.
|
||||
main_qkvg_width = 3584
|
||||
projected_qk = torch.randn(
|
||||
num_tokens, (HQ + 1) * D, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
num_tokens, main_qkvg_width + (HQ + 1) * D, dtype=torch.bfloat16, device=device
|
||||
)[:, main_qkvg_width:]
|
||||
q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
|
||||
k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
|
||||
|
||||
|
||||
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
QWEN4_EXP_SM103_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
# GDN fused QKVZ projection, TP=4.
|
||||
(4096, 2560): {
|
||||
1: SkinnyGemmConfig(1, 64, 4, k_unroll=4),
|
||||
@@ -43,18 +43,11 @@ QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
|
||||
16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4),
|
||||
2: SkinnyGemmConfig(2, 64, 2, k_unroll=2),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4),
|
||||
2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4 (GB300).
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# Shared-expert fused gate/up projection, TP=4.
|
||||
(320, 2560): {
|
||||
@@ -69,13 +62,6 @@ QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
1: SkinnyGemmConfig(1, 64, 4, k_unroll=2),
|
||||
2: SkinnyGemmConfig(2, 32, 4, k_unroll=2),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 128, 1, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 128, 1, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 128, 2, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4),
|
||||
},
|
||||
}
|
||||
|
||||
# H200 plans selected by exhaustive CUDA graph replay measurements over
|
||||
@@ -88,6 +74,12 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4.
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=6, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, k_unroll=3, vector_width=2, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4),
|
||||
},
|
||||
# GDN and QSA output projections, TP=4.
|
||||
(2560, 1536): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=1536),
|
||||
@@ -102,18 +94,6 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 4, k_unroll=5),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 2, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 128, 2, vector_width=4, static_k=2560),
|
||||
},
|
||||
# Shared-expert fused gate/up projection, TP=4.
|
||||
(320, 2560): {
|
||||
1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=2560),
|
||||
@@ -127,13 +107,6 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
1: SkinnyGemmConfig(1, 64, 2, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 2, vector_width=2, static_k=2560),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=5),
|
||||
2: SkinnyGemmConfig(2, 256, 3, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 256, 3, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
|
||||
},
|
||||
# Final HC down projection, replicated in a TP=4 deployment.
|
||||
(320, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, static_k=10240),
|
||||
@@ -169,30 +142,17 @@ QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
|
||||
8: SkinnyGemmConfig(8, 256, 1),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN and QSA output projections, TP=4.
|
||||
(2560, 1536): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=1536),
|
||||
2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=1536),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=2, static_k=1536),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4.
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN fused QKVZ projection, TP=4.
|
||||
(4096, 2560): {
|
||||
@@ -207,7 +167,7 @@ QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
}
|
||||
|
||||
|
||||
# DGX Spark (GB10) plans, TP=1.
|
||||
# DGX Spark (GB10) plans, TP=1 unless noted.
|
||||
QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
# Shared-expert gate.
|
||||
(1, 2560): {
|
||||
@@ -225,14 +185,6 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 64, 1),
|
||||
},
|
||||
# HC merged down/injection projection.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, vector_width=4, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=4, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 128, 4, k_unroll=4, static_k=10240),
|
||||
16: SkinnyGemmConfig(16, 128, 4, k_unroll=4, static_k=10240),
|
||||
},
|
||||
# Router.
|
||||
(512, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560),
|
||||
@@ -241,7 +193,7 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 64, 2, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 64, 2, k_unroll=2, static_k=2560),
|
||||
},
|
||||
# QSA indexer Q/K projection.
|
||||
# Shared-expert fused gate/up, TP=2; retain after indexer fusion.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 256, 1, vector_width=2, static_k=2560),
|
||||
@@ -280,13 +232,13 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=320),
|
||||
8: SkinnyGemmConfig(8, 32, 1, vector_width=2, static_k=320),
|
||||
},
|
||||
# QSA fused QKV/gate projection.
|
||||
(13312, 2560): {
|
||||
# QSA fused QKV/gate + replicated indexer Q/K projection.
|
||||
(13952, 2560): {
|
||||
1: SkinnyGemmConfig(1, 32, 1, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 32, 1, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 32, 1, k_unroll=2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 32, 1, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN fused QKVZ projection.
|
||||
(16384, 2560): {
|
||||
@@ -307,32 +259,14 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
}
|
||||
|
||||
|
||||
def _is_sm121() -> bool:
|
||||
return current_platform.is_device_capability((12, 1))
|
||||
|
||||
|
||||
def _is_sm100() -> bool:
|
||||
return current_platform.is_device_capability((10, 0))
|
||||
|
||||
|
||||
def _is_sm103() -> bool:
|
||||
return current_platform.is_device_capability((10, 3))
|
||||
|
||||
|
||||
def _is_sm90() -> bool:
|
||||
return current_platform.is_device_capability((9, 0))
|
||||
|
||||
|
||||
def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
|
||||
if _is_sm103():
|
||||
return QWEN4_EXP_GEMM_PLANS
|
||||
if _is_sm100():
|
||||
return QWEN4_EXP_SM100_GEMM_PLANS
|
||||
if _is_sm90():
|
||||
return QWEN4_EXP_SM90_GEMM_PLANS
|
||||
if _is_sm121():
|
||||
return QWEN4_EXP_SM121_GEMM_PLANS
|
||||
return {}
|
||||
QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY: dict[
|
||||
tuple[int, int] | None, dict[tuple[int, int], dict[int, SkinnyGemmConfig]]
|
||||
] = {
|
||||
(10, 3): QWEN4_EXP_SM103_GEMM_PLANS,
|
||||
(10, 0): QWEN4_EXP_SM100_GEMM_PLANS,
|
||||
(9, 0): QWEN4_EXP_SM90_GEMM_PLANS,
|
||||
(12, 1): QWEN4_EXP_SM121_GEMM_PLANS,
|
||||
}
|
||||
|
||||
|
||||
def _is_packed_row_major(tensor: torch.Tensor) -> bool:
|
||||
@@ -376,7 +310,10 @@ class Qwen4ExpLowLatencyEmbeddingMethod(
|
||||
|
||||
|
||||
def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
plan = _gemm_plans().get((weight.shape[0], weight.shape[1]))
|
||||
plans = QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
current_platform.get_device_capability(), {}
|
||||
)
|
||||
plan = plans.get((weight.shape[0], weight.shape[1]))
|
||||
config = None if plan is None else plan.get(x.shape[0])
|
||||
if (
|
||||
config is not None
|
||||
@@ -404,7 +341,9 @@ def enable_qwen4_exp_low_latency_gemm(
|
||||
module: nn.Module,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
plans = _gemm_plans()
|
||||
plans = QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
current_platform.get_device_capability(), {}
|
||||
)
|
||||
if dtype != torch.bfloat16 or not plans:
|
||||
return
|
||||
if not shape_dynamic_skinny_gemm.is_available():
|
||||
|
||||
@@ -622,7 +622,7 @@ class Qwen4ExpForCausalLM(
|
||||
IsHybrid,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj", "indexer.index_qk_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
@@ -641,6 +641,12 @@ class Qwen4ExpForCausalLM(
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
|
||||
if vllm_config.lora_config is not None:
|
||||
# LoRA does not support the merged QKV/indexer projection, so its
|
||||
# packed mapping must keep the indexer separate.
|
||||
self.packed_modules_mapping = self.packed_modules_mapping | {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
}
|
||||
self.vllm_config = vllm_config
|
||||
self.model_config = vllm_config.model_config
|
||||
self.quant_config = vllm_config.quant_config
|
||||
@@ -866,6 +872,7 @@ class Qwen4ExpForConditionalGeneration(
|
||||
requires_raw_input_tokens = True
|
||||
|
||||
packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
|
||||
"qkv_proj": Qwen4ExpForCausalLM.packed_modules_mapping["qkv_proj"],
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"input_mix_weight_down_block_inject": [
|
||||
"input_mix_weight_down",
|
||||
@@ -883,6 +890,12 @@ class Qwen4ExpForConditionalGeneration(
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen4ExpConfig = vllm_config.model_config.hf_config
|
||||
if vllm_config.lora_config is not None:
|
||||
# LoRA does not support the merged QKV/indexer projection, so its
|
||||
# packed mapping must keep the indexer separate.
|
||||
self.packed_modules_mapping = self.packed_modules_mapping | {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
}
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
if multimodal_config is None:
|
||||
|
||||
@@ -154,6 +154,7 @@ def _qsa_prepare_kernel(
|
||||
MROPE_H: tl.constexpr,
|
||||
MROPE_W: tl.constexpr,
|
||||
main_qkv_ptr,
|
||||
main_qkv_stride_token,
|
||||
main_q_norm_weight_ptr,
|
||||
main_k_norm_weight_ptr,
|
||||
main_eps,
|
||||
@@ -178,12 +179,12 @@ def _qsa_prepare_kernel(
|
||||
# indexer work. RoPE covers the first D // 2 dims, the width of the
|
||||
# shared cos/sin table.
|
||||
main_pid = pid - num_index_work
|
||||
token = main_pid // (MAIN_HQ + MAIN_HK)
|
||||
token = (main_pid // (MAIN_HQ + MAIN_HK)).to(tl.int64)
|
||||
head = main_pid % (MAIN_HQ + MAIN_HK)
|
||||
HALF: tl.constexpr = D // 4
|
||||
dims = tl.arange(0, MAIN_D)
|
||||
rot = tl.arange(0, HALF)
|
||||
row = main_qkv_ptr + token * (2 * (MAIN_HQ + MAIN_HK) * MAIN_D)
|
||||
row = main_qkv_ptr + token * main_qkv_stride_token
|
||||
is_k = head >= MAIN_HQ
|
||||
kv_head = head - MAIN_HQ
|
||||
if is_k:
|
||||
@@ -241,7 +242,7 @@ def _qsa_prepare_kernel(
|
||||
num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q)
|
||||
token_tile = q_pid // num_head_tiles
|
||||
head_tile = q_pid % num_head_tiles
|
||||
tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)
|
||||
tokens = (token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)).to(tl.int64)
|
||||
heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q)
|
||||
valid_tokens = tokens < num_tokens
|
||||
valid_heads = heads < HQ
|
||||
@@ -342,7 +343,8 @@ def _qsa_prepare_kernel(
|
||||
& (source_tokens < num_tokens)
|
||||
)
|
||||
current_base = (
|
||||
k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token
|
||||
k_ptr
|
||||
+ tl.maximum(source_tokens, 0).to(tl.int64)[:, None] * k_stride_token
|
||||
)
|
||||
cached_base = (
|
||||
state_cache_ptr
|
||||
@@ -449,7 +451,7 @@ def _qsa_prepare_kernel(
|
||||
# One CTA per request commits only the suffix retained by the ring.
|
||||
num_state_rows = tl.minimum(query_len, STATE_SIZE)
|
||||
for state_offset in tl.range(0, num_state_rows):
|
||||
token = query_end - num_state_rows + state_offset
|
||||
token = (query_end - num_state_rows + state_offset).to(tl.int64)
|
||||
valid_token = (
|
||||
(token >= query_start) & (token < query_end) & (token < num_tokens)
|
||||
)
|
||||
@@ -567,7 +569,8 @@ def qsa_prepare(
|
||||
section = mrope_section if mrope_section is not None else (0, 0, 0)
|
||||
assert len(section) == 3
|
||||
qkv_width = 2 * (num_main_q_heads + num_main_kv_heads) * main_head_dim
|
||||
assert main_qkv.shape == (num_tokens, qkv_width) and main_qkv.is_contiguous()
|
||||
assert main_qkv.shape == (num_tokens, qkv_width)
|
||||
assert main_qkv.stride(-1) == 1
|
||||
assert main_slot_mapping.shape == (num_tokens,)
|
||||
|
||||
if num_tokens <= 4096:
|
||||
@@ -622,6 +625,7 @@ def qsa_prepare(
|
||||
MROPE_H=section[1],
|
||||
MROPE_W=section[2],
|
||||
main_qkv_ptr=main_qkv,
|
||||
main_qkv_stride_token=main_qkv.stride(0),
|
||||
main_q_norm_weight_ptr=main_q_norm_weight,
|
||||
main_k_norm_weight_ptr=main_k_norm_weight,
|
||||
main_eps=main_eps,
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import ClassVar, cast
|
||||
|
||||
import torch
|
||||
@@ -20,10 +21,15 @@ from vllm.model_executor.layers.attention.attention import (
|
||||
)
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
|
||||
from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.utils.torch_utils import (
|
||||
@@ -53,6 +59,72 @@ from . import model
|
||||
from .indexer_qsa import QSAIndexer
|
||||
|
||||
|
||||
class Qwen4ExpQSAQKVIndexerLinear(MergedColumnParallelLinear):
|
||||
"""Pack sharded Q/gate, K, V and replicated indexer Q/K in one GEMM."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
qkv_proj: QKVParallelLinear,
|
||||
index_size: int,
|
||||
quant_config: QuantizationConfig | None,
|
||||
) -> None:
|
||||
self.num_kv_head_replicas = qkv_proj.num_kv_head_replicas
|
||||
super().__init__(
|
||||
input_size=qkv_proj.input_size,
|
||||
output_sizes=[*qkv_proj.output_sizes, index_size * qkv_proj.tp_size],
|
||||
bias=False,
|
||||
params_dtype=qkv_proj.params_dtype,
|
||||
quant_config=quant_config,
|
||||
prefix=qkv_proj.prefix,
|
||||
)
|
||||
|
||||
def _load_shard(
|
||||
self,
|
||||
loader: Callable[..., None],
|
||||
param: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: int | tuple[int, ...] | None,
|
||||
) -> None:
|
||||
tp_rank = self.tp_rank
|
||||
param_tp_rank = getattr(param, "tp_rank", None)
|
||||
if loaded_shard_id == 3:
|
||||
shard_rank = 0
|
||||
elif loaded_shard_id in (1, 2):
|
||||
shard_rank = tp_rank // self.num_kv_head_replicas
|
||||
else:
|
||||
shard_rank = tp_rank
|
||||
self.tp_rank = shard_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = shard_rank
|
||||
try:
|
||||
loader(param, loaded_weight, loaded_shard_id)
|
||||
finally:
|
||||
self.tp_rank = tp_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = param_tp_rank
|
||||
|
||||
def weight_loader(self, param, loaded_weight, loaded_shard_id=None) -> None:
|
||||
self._load_shard(super().weight_loader, param, loaded_weight, loaded_shard_id)
|
||||
|
||||
def weight_loader_v2(self, param, loaded_weight, loaded_shard_id=None) -> None:
|
||||
self._load_shard(
|
||||
super().weight_loader_v2, param, loaded_weight, loaded_shard_id
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self, weights: Iterable[tuple[str, torch.Tensor]]
|
||||
) -> Iterable[str]:
|
||||
def remap_shards():
|
||||
for name, weight in weights:
|
||||
shard_id = getattr(weight, "shard_id", None)
|
||||
if isinstance(shard_id, str):
|
||||
weight = weight.detach()
|
||||
weight.shard_id = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
yield name, weight
|
||||
|
||||
return super().load_weights(remap_shards())
|
||||
|
||||
|
||||
class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
|
||||
"""Flash metadata supporting uniform decode and target-verify graphs."""
|
||||
|
||||
@@ -411,6 +483,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
self.use_fused_qsa_prepare = (
|
||||
self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer
|
||||
)
|
||||
self.fuse_indexer_projection = vllm_config.lora_config is None
|
||||
if self.fuse_indexer_projection:
|
||||
self.index_qk_size = self.indexer.index_qk_proj.output_size
|
||||
self.qkv_proj = Qwen4ExpQSAQKVIndexerLinear(
|
||||
self.qkv_proj,
|
||||
self.index_qk_size,
|
||||
model.without_modelopt_fp4(quant_config),
|
||||
)
|
||||
del self.indexer.index_qk_proj
|
||||
|
||||
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
# PACKED selection buffer: the trailing column holds each row's
|
||||
# valid-entry count (written by the expand kernel) — never a token
|
||||
@@ -432,6 +514,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
raise ValueError(f"Duplicate layer name: {self.layer_name}")
|
||||
static_context[self.layer_name] = self
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
mapper = None
|
||||
if self.fuse_indexer_projection:
|
||||
mapper = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
"indexer.index_qk_proj.": ("qkv_proj.", 3),
|
||||
}
|
||||
)
|
||||
return AutoWeightsLoader(self).load_weights(weights, mapper=mapper)
|
||||
|
||||
def get_attn_backend(self) -> type[AttentionBackend]:
|
||||
return self.attn_backend
|
||||
|
||||
@@ -518,6 +610,12 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
if self.fuse_indexer_projection:
|
||||
qkv, projected_qk = qkv.split(
|
||||
[2 * self.q_size + 2 * self.kv_size, self.index_qk_size], dim=-1
|
||||
)
|
||||
else:
|
||||
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
|
||||
num_tokens = hidden_states.shape[0]
|
||||
if not self.use_fused_qsa_prepare:
|
||||
q, k, v, gate = self._project_qkv_gate(qkv, positions)
|
||||
@@ -529,8 +627,6 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
# Norm/RoPE/gate and the K/V cache write happen inside _run_qsa.
|
||||
query = key = value = gate = None
|
||||
attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim)
|
||||
# Keep the index projection outside the eager break.
|
||||
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
|
||||
self._run_qsa(
|
||||
projected_qk,
|
||||
positions,
|
||||
|
||||
Reference in New Issue
Block a user