diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 7cdaf2b495..7ca207c8e4 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -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( diff --git a/tests/models/qwen4_exp/test_qsa_prepare.py b/tests/models/qwen4_exp/test_qsa_prepare.py index 14b836c132..e56a50b6c9 100644 --- a/tests/models/qwen4_exp/test_qsa_prepare.py +++ b/tests/models/qwen4_exp/test_qsa_prepare.py @@ -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 diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index fa01c75698..80bd4d2ddc 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -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(): diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 73fe2cc468..2de883cc47 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -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: diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py index 6e97a89781..a26b699657 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py @@ -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, diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index d00cd186b0..6877189768 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -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,