From 02f2d4c8f139ca7cd3811ee9e38a78ea12de9af5 Mon Sep 17 00:00:00 2001 From: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Date: Mon, 31 Aug 2026 20:03:56 +0800 Subject: [PATCH 1/4] [Kernel][Qwen] Fuse Qwen4Exp decode HyperConnection Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Assisted-by: OpenAI Codex --- tests/models/qwen4_exp/test_hc_ops.py | 29 +++ .../qwen4_exp/nvidia/hyperconnection.py | 34 +++- .../nvidia/ops/_hc_silu_up_gate_mix.py | 145 ++++++++++++++ .../nvidia/ops/hc_silu_up_gate_mix.py | 185 ++++++++++++++++++ 4 files changed, 387 insertions(+), 6 deletions(-) create mode 100644 vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py diff --git a/tests/models/qwen4_exp/test_hc_ops.py b/tests/models/qwen4_exp/test_hc_ops.py index 4af4e5573232..3c75715c6b9c 100644 --- a/tests/models/qwen4_exp/test_hc_ops.py +++ b/tests/models/qwen4_exp/test_hc_ops.py @@ -9,6 +9,11 @@ hc_combine, hc_combine_norm, hc_gate_mix, + hc_silu, +) +from vllm.models.qwen4_exp.nvidia.ops.hc_silu_up_gate_mix import ( + HCSiluUpGateMixOp, + hc_silu_up_gate_mix, ) from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON @@ -52,6 +57,30 @@ def test_hc_gate_mix() -> None: torch.testing.assert_close(actual, expected.to(torch.bfloat16)) +@pytest.mark.skipif( + not current_platform.is_device_capability((9, 0)), + reason="fused HC up-projection requires SM90", +) +@pytest.mark.parametrize("num_tokens", [1, 2]) +def test_hc_silu_up_gate_mix(num_tokens: int) -> None: + if not HCSiluUpGateMixOp.is_supported(torch.bfloat16): + pytest.skip("CuTeDSL is not available") + + torch.manual_seed(0) + # The model obtains lora by splitting the padded down projection. + lora_storage = torch.randn(num_tokens, 336, dtype=torch.bfloat16, device="cuda") + lora = lora_storage[:, :320] + weight = torch.randn(HYPER_HIDDEN_SIZE, 320, dtype=torch.bfloat16, device="cuda") + weight /= 320**0.5 + x = torch.randn(num_tokens, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + + actual = hc_silu_up_gate_mix(lora, weight, x) + gate = torch.nn.functional.linear(hc_silu(lora, HC), weight) + expected = hc_gate_mix(x, gate, HC) + + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-3) + + def test_hc_combine() -> None: torch.manual_seed(0) block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index 8ca503d214c8..7c79b946baf9 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -28,6 +28,7 @@ from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, ReplicatedLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.models.utils import maybe_prefix @@ -42,6 +43,7 @@ hc_gate_mix, hc_silu, ) +from .ops.hc_silu_up_gate_mix import HCSiluUpGateMixOp, hc_silu_up_gate_mix # --------------------------------------------------------------------------- @@ -123,6 +125,30 @@ def __init__( prefix=maybe_prefix(prefix, "input_mix_weight_up"), return_bias=False, ) + self.use_fused_hc_up = ( + config.hc_count == 4 + and config.hidden_size == 2560 + and config.hc_lowrank == 320 + and type(self.input_mix_weight_up.quant_method) is UnquantizedLinearMethod + and HCSiluUpGateMixOp.is_supported(config.params_dtype) + ) + if self.use_fused_hc_up: + HCSiluUpGateMixOp.initialize() + + def _up_and_gate_mix( + self, + lora: torch.Tensor, + xn: torch.Tensor, + ) -> torch.Tensor: + if self.use_fused_hc_up: + return hc_silu_up_gate_mix( + lora, + self.input_mix_weight_up.weight, + xn, + ) + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) + return hc_gate_mix(xn, gate, self.hc_count) def mix( self, hidden_states: torch.Tensor @@ -143,9 +169,7 @@ def mix( lora = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] - block_input = hc_gate_mix(xn, gate, self.hc_count) + block_input = self._up_and_gate_mix(lora, xn) return hidden_states, block_input, injection @@ -179,9 +203,7 @@ def combine_and_mix( lora = self.input_mix_weight_down(xn) injection = None - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) # [M, D] - block_input = hc_gate_mix(xn, gate, self.hc_count) + block_input = self._up_and_gate_mix(lora, xn) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py b/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py new file mode 100644 index 000000000000..3ae6f0cdf19d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream + +_HC = 4 +_K = 320 +_BLOCK_SIZE = 32 +_OUTPUTS_PER_BLOCK = 2 +_VECTOR_WIDTH = 2 + + +class HCSiluUpGateMixKernel: + """Fused Qwen4Exp decode HC up-projection and gated stream reduction.""" + + def __init__(self) -> None: + self.element_type = cutlass.BFloat16 + + @cute.jit + def __call__( + self, + g_lora: cute.Tensor, + g_weight: cute.Tensor, + g_x: cute.Tensor, + g_out: cute.Tensor, + stream: CUstream, + ) -> None: + hidden_size = cute.size(g_out, mode=[1]) + copy_lora = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=_VECTOR_WIDTH * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_weight = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=_VECTOR_WIDTH * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + g_lora, + g_weight, + g_x, + g_out, + hidden_size, + copy_lora, + copy_weight, + ).launch( + grid=[cute.ceil_div(hidden_size, _OUTPUTS_PER_BLOCK), 1, 1], + block=[_BLOCK_SIZE, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + g_lora: cute.Tensor, + g_weight: cute.Tensor, + g_x: cute.Tensor, + g_out: cute.Tensor, + hidden_size: cutlass.Int32, + copy_lora: cute.CopyAtom, + copy_weight: cute.CopyAtom, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + + acc = cute.make_rmem_tensor( + cute.make_layout( + (_OUTPUTS_PER_BLOCK, _HC), + stride=(_HC, 1), + ), + cutlass.Float32, + ) + acc.fill(0.0) + + lora_vec = cute.logical_divide(g_lora, (None, _VECTOR_WIDTH)) + weight_vec = cute.logical_divide(g_weight, (None, _VECTOR_WIDTH)) + lora_tiles = cute.logical_divide(lora_vec, (None, (None, _BLOCK_SIZE))) + weight_tiles = cute.logical_divide(weight_vec, (None, (None, _BLOCK_SIZE))) + thread_lora = lora_tiles[None, (None, (tidx, None))] + lora_regs = cute.make_rmem_tensor( + cute.make_layout((_VECTOR_WIDTH,), stride=(1,)), + self.element_type, + ) + weight_regs = cute.make_rmem_tensor( + cute.make_layout( + (_OUTPUTS_PER_BLOCK, _HC, _VECTOR_WIDTH), + stride=(_HC * _VECTOR_WIDTH, _VECTOR_WIDTH, 1), + ), + self.element_type, + ) + + hidden_base = block_idx * _OUTPUTS_PER_BLOCK + for k_tile in cutlass.range_constexpr(_K // (_BLOCK_SIZE * _VECTOR_WIDTH)): + cute.copy(copy_lora, thread_lora[0, None, k_tile], lora_regs) + for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): + for hc_stream in cutlass.range_constexpr(_HC): + weight_row = hc_stream * hidden_size + hidden_base + output + thread_weight = weight_tiles[weight_row, (None, (tidx, None))] + cute.copy( + copy_weight, + thread_weight[None, k_tile], + weight_regs[output, hc_stream, None], + ) + + raw_lora = lora_regs.load().to(cutlass.Float32) / _HC + activated_lora = (raw_lora / (1.0 + cute.exp(-raw_lora, fastmath=True))).to( + self.element_type + ) + activated_lora = activated_lora.to(cutlass.Float32) + weights = weight_regs.load().to(cutlass.Float32) + for vector_lane in cutlass.range_constexpr(_VECTOR_WIDTH): + for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): + for hc_stream in cutlass.range_constexpr(_HC): + acc[output, hc_stream] += ( + activated_lora[vector_lane] + * weights[output, hc_stream, vector_lane] + ) + + for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): + for hc_stream in cutlass.range_constexpr(_HC): + acc[output, hc_stream] = cute.arch.warp_reduction_sum( + acc[output, hc_stream] + ) + + if tidx == 0: + for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): + hidden_idx = hidden_base + output + if hidden_idx < hidden_size: + mixed = cutlass.Float32(0.0) + for hc_stream in cutlass.range_constexpr(_HC): + x = g_x[0, hc_stream * hidden_size + hidden_idx].to( + cutlass.Float32 + ) + mixed += x / ( + 1.0 + cute.exp(-acc[output, hc_stream], fastmath=True) + ) + g_out[0, hidden_idx] = (mixed / _HC).to(self.element_type) diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py b/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py new file mode 100644 index 000000000000..bcd87e938af3 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any, ClassVar + +import torch + +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import current_stream, direct_register_custom_op + +_HC = 4 +_HIDDEN_SIZE = 2560 +_HYPER_HIDDEN_SIZE = _HC * _HIDDEN_SIZE +_LORA_RANK = 320 + + +class HCSiluUpGateMixOp: + """Process-local compiled fused HyperConnection decode operation.""" + + _instance: ClassVar[HCSiluUpGateMixOp | None] = None + + @classmethod + def initialize(cls) -> HCSiluUpGateMixOp: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + @staticmethod + def is_supported(dtype: torch.dtype) -> bool: + if dtype != torch.bfloat16 or not current_platform.is_device_capability((9, 0)): + return False + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + except ImportError: + return False + return True + + def __init__(self) -> None: + self._compiled: Any | None = None + register_cutedsl_warmup_provider(self) + + def _compile(self) -> None: + import cutlass + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + from quack.compile_utils import make_fake_tensor + + from ._hc_silu_up_gate_mix import HCSiluUpGateMixKernel + + lora = make_fake_tensor( + cutlass.BFloat16, + (1, _LORA_RANK), + divisibility=2, + ) + weight = make_fake_tensor( + cutlass.BFloat16, + (_HYPER_HIDDEN_SIZE, _LORA_RANK), + divisibility=2, + ) + x = make_fake_tensor( + cutlass.BFloat16, + (1, _HYPER_HIDDEN_SIZE), + divisibility=1, + ) + out = make_fake_tensor( + cutlass.BFloat16, + (1, _HIDDEN_SIZE), + divisibility=1, + ) + self._compiled = cute.compile( + HCSiluUpGateMixKernel(), + lora, + weight, + x, + out, + CUstream(current_stream().cuda_stream), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", + ) + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + return ( + CuTeDSLCompileUnit( + name="Qwen4Exp fused HC SiLU/up-projection/gate-mix", + key=("qwen4-exp-hc-silu-up-gate-mix", torch.bfloat16), + compile=self._compile, + ), + ) + + def __call__( + self, + lora: torch.Tensor, + weight: torch.Tensor, + x: torch.Tensor, + ) -> torch.Tensor: + self._validate_inputs(lora, weight, x) + if self._compiled is None: + self._compile() + compiled = self._compiled + assert compiled is not None + out = torch.empty( + (1, _HIDDEN_SIZE), + dtype=torch.bfloat16, + device=lora.device, + ) + from cuda.bindings.driver import CUstream + + compiled( + lora, + weight, + x, + out, + CUstream(current_stream().cuda_stream), + ) + return out + + @staticmethod + def _validate_inputs( + lora: torch.Tensor, + weight: torch.Tensor, + x: torch.Tensor, + ) -> None: + if lora.shape != (1, _LORA_RANK): + raise ValueError(f"lora must have shape [1, {_LORA_RANK}].") + if weight.shape != (_HYPER_HIDDEN_SIZE, _LORA_RANK): + raise ValueError( + f"weight must have shape [{_HYPER_HIDDEN_SIZE}, {_LORA_RANK}]." + ) + if x.shape != (1, _HYPER_HIDDEN_SIZE): + raise ValueError(f"x must have shape [1, {_HYPER_HIDDEN_SIZE}].") + + tensors = (lora, weight, x) + if any(tensor.dtype != torch.bfloat16 for tensor in tensors): + raise ValueError("All inputs must use torch.bfloat16.") + if any(tensor.device != lora.device for tensor in tensors): + raise ValueError("All inputs must be on the same device.") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("All inputs must be contiguous.") + + +def _hc_silu_up_gate_mix( + lora: torch.Tensor, + weight: torch.Tensor, + x: torch.Tensor, +) -> torch.Tensor: + if lora.shape[0] == 1: + return HCSiluUpGateMixOp.initialize()(lora, weight, x) + + from .hc import hc_gate_mix, hc_silu + + gate = torch.nn.functional.linear(hc_silu(lora, _HC), weight) + return hc_gate_mix(x, gate, _HC) + + +def _hc_silu_up_gate_mix_fake( + lora: torch.Tensor, + weight: torch.Tensor, + x: torch.Tensor, +) -> torch.Tensor: + del weight, x + return lora.new_empty((lora.shape[0], _HIDDEN_SIZE)) + + +direct_register_custom_op( + op_name="qwen4_exp_hc_silu_up_gate_mix", + op_func=_hc_silu_up_gate_mix, + fake_impl=_hc_silu_up_gate_mix_fake, +) + + +def hc_silu_up_gate_mix( + lora: torch.Tensor, + weight: torch.Tensor, + x: torch.Tensor, +) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_silu_up_gate_mix(lora, weight, x) + + +__all__ = ["HCSiluUpGateMixOp", "hc_silu_up_gate_mix"] From ce4f884951788d1bf90cc3b26fc82a7d87fc437f Mon Sep 17 00:00:00 2001 From: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:36:38 +0800 Subject: [PATCH 2/4] [Kernel][Qwen] Enable Hopper decode LL-GEMM Use the measured Qwen4Exp low-latency GEMM plans for M=1 on SM90 while retaining the standard linear fallback for larger batches. Assisted-by: OpenAI Codex Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> --- tests/kernels/test_bf16_skinny_gemm.py | 32 ++++++++++++++++++- .../qwen4_exp/nvidia/low_latency_gemm.py | 27 +++++++++++++--- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index c20823138834..9021b1d640b0 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM90/SM100/SM103 selectors.""" +"""Tests for BF16 skinny GEMMs and model-specific selectors.""" from pathlib import Path from types import SimpleNamespace @@ -17,6 +17,7 @@ from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS +from vllm.models.qwen4_exp.nvidia import low_latency_gemm as qwen4_exp_gemm # Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is # the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16). @@ -528,6 +529,35 @@ def test_low_latency_table_capability_routing( assert k3_gemm._low_latency_table() is None +def test_qwen4_exp_hopper_plans_are_decode_only() -> None: + plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS + + assert plans.keys() == qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS.keys() + assert all(set(shape_plans) == {1} for shape_plans in plans.values()) + + +@pytest.mark.parametrize( + "capability,expected_plans", + [ + ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS), + ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS), + ((8, 0), {}), + ], +) +def test_qwen4_exp_gemm_capability_routing( + monkeypatch: pytest.MonkeyPatch, + capability: tuple[int, int], + expected_plans: dict[tuple[int, int], dict[int, SkinnyGemmConfig]], +) -> None: + monkeypatch.setattr( + qwen4_exp_gemm.current_platform, + "is_device_capability", + lambda target: capability == target, + ) + + assert qwen4_exp_gemm._gemm_plans() == expected_plans + + def test_installation_is_shape_specific_and_unquantized( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index 43b7a76bc5d8..4f963f6bc082 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Qwen4Exp decode GEMM selection on Blackwell. +"""Qwen4Exp decode GEMM selection on Hopper and Blackwell. Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count. Plans contain measured CUDA graph capture sizes; other token counts use the @@ -78,11 +78,29 @@ }, } +# H200 plans measured under CUDA graph replay. Only M=1 is enabled so larger +# batches retain the standard linear implementation and its GEMM heuristics. +QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { + shape: {1: plans[1]} for shape, plans in QWEN4_EXP_GEMM_PLANS.items() +} + 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_sm90(): + return QWEN4_EXP_SM90_GEMM_PLANS + return {} + + def _is_packed_row_major(tensor: torch.Tensor) -> bool: return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) @@ -124,7 +142,7 @@ class Qwen4ExpLowLatencyEmbeddingMethod( def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + plan = _gemm_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 @@ -152,7 +170,8 @@ def enable_qwen4_exp_low_latency_gemm( module: nn.Module, dtype: torch.dtype, ) -> None: - if dtype != torch.bfloat16 or not _is_sm103(): + plans = _gemm_plans() + if dtype != torch.bfloat16 or not plans: return if not shape_dynamic_skinny_gemm.is_available(): return @@ -172,7 +191,7 @@ def enable_qwen4_exp_low_latency_gemm( weight = getattr(child, "weight", None) if weight is None or weight.dim() != 2: continue - plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + plan = plans.get((weight.shape[0], weight.shape[1])) if plan is None: continue if is_linear: From 9bed9603c55e851c639288499cbc84e3f62ae4a7 Mon Sep 17 00:00:00 2001 From: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:04:30 +0800 Subject: [PATCH 3/4] [Kernel][Qwen] Drop Hopper-specific HC fusion Keep the HyperConnection path on the standard linear implementation so Hopper decode optimization can use the shared LL-GEMM infrastructure instead of a model- and shape-specific fused kernel.\n\nAssisted-by: OpenAI Codex Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> --- tests/models/qwen4_exp/test_hc_ops.py | 29 --- .../qwen4_exp/nvidia/hyperconnection.py | 34 +--- .../nvidia/ops/_hc_silu_up_gate_mix.py | 145 -------------- .../nvidia/ops/hc_silu_up_gate_mix.py | 185 ------------------ 4 files changed, 6 insertions(+), 387 deletions(-) delete mode 100644 vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py delete mode 100644 vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py diff --git a/tests/models/qwen4_exp/test_hc_ops.py b/tests/models/qwen4_exp/test_hc_ops.py index 3c75715c6b9c..4af4e5573232 100644 --- a/tests/models/qwen4_exp/test_hc_ops.py +++ b/tests/models/qwen4_exp/test_hc_ops.py @@ -9,11 +9,6 @@ hc_combine, hc_combine_norm, hc_gate_mix, - hc_silu, -) -from vllm.models.qwen4_exp.nvidia.ops.hc_silu_up_gate_mix import ( - HCSiluUpGateMixOp, - hc_silu_up_gate_mix, ) from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON @@ -57,30 +52,6 @@ def test_hc_gate_mix() -> None: torch.testing.assert_close(actual, expected.to(torch.bfloat16)) -@pytest.mark.skipif( - not current_platform.is_device_capability((9, 0)), - reason="fused HC up-projection requires SM90", -) -@pytest.mark.parametrize("num_tokens", [1, 2]) -def test_hc_silu_up_gate_mix(num_tokens: int) -> None: - if not HCSiluUpGateMixOp.is_supported(torch.bfloat16): - pytest.skip("CuTeDSL is not available") - - torch.manual_seed(0) - # The model obtains lora by splitting the padded down projection. - lora_storage = torch.randn(num_tokens, 336, dtype=torch.bfloat16, device="cuda") - lora = lora_storage[:, :320] - weight = torch.randn(HYPER_HIDDEN_SIZE, 320, dtype=torch.bfloat16, device="cuda") - weight /= 320**0.5 - x = torch.randn(num_tokens, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") - - actual = hc_silu_up_gate_mix(lora, weight, x) - gate = torch.nn.functional.linear(hc_silu(lora, HC), weight) - expected = hc_gate_mix(x, gate, HC) - - torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-3) - - def test_hc_combine() -> None: torch.manual_seed(0) block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index 7c79b946baf9..8ca503d214c8 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -28,7 +28,6 @@ from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, ReplicatedLinear, - UnquantizedLinearMethod, ) from vllm.model_executor.models.utils import maybe_prefix @@ -43,7 +42,6 @@ hc_gate_mix, hc_silu, ) -from .ops.hc_silu_up_gate_mix import HCSiluUpGateMixOp, hc_silu_up_gate_mix # --------------------------------------------------------------------------- @@ -125,30 +123,6 @@ def __init__( prefix=maybe_prefix(prefix, "input_mix_weight_up"), return_bias=False, ) - self.use_fused_hc_up = ( - config.hc_count == 4 - and config.hidden_size == 2560 - and config.hc_lowrank == 320 - and type(self.input_mix_weight_up.quant_method) is UnquantizedLinearMethod - and HCSiluUpGateMixOp.is_supported(config.params_dtype) - ) - if self.use_fused_hc_up: - HCSiluUpGateMixOp.initialize() - - def _up_and_gate_mix( - self, - lora: torch.Tensor, - xn: torch.Tensor, - ) -> torch.Tensor: - if self.use_fused_hc_up: - return hc_silu_up_gate_mix( - lora, - self.input_mix_weight_up.weight, - xn, - ) - lora = hc_silu(lora, self.hc_count) - gate = self.input_mix_weight_up(lora) - return hc_gate_mix(xn, gate, self.hc_count) def mix( self, hidden_states: torch.Tensor @@ -169,7 +143,9 @@ def mix( lora = self.input_mix_weight_down(xn) injection = None - block_input = self._up_and_gate_mix(lora, xn) + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection @@ -203,7 +179,9 @@ def combine_and_mix( lora = self.input_mix_weight_down(xn) injection = None - block_input = self._up_and_gate_mix(lora, xn) + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) return hidden_states, block_input, injection diff --git a/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py b/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py deleted file mode 100644 index 3ae6f0cdf19d..000000000000 --- a/vllm/models/qwen4_exp/nvidia/ops/_hc_silu_up_gate_mix.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import cutlass -import cutlass.cute as cute -from cuda.bindings.driver import CUstream - -_HC = 4 -_K = 320 -_BLOCK_SIZE = 32 -_OUTPUTS_PER_BLOCK = 2 -_VECTOR_WIDTH = 2 - - -class HCSiluUpGateMixKernel: - """Fused Qwen4Exp decode HC up-projection and gated stream reduction.""" - - def __init__(self) -> None: - self.element_type = cutlass.BFloat16 - - @cute.jit - def __call__( - self, - g_lora: cute.Tensor, - g_weight: cute.Tensor, - g_x: cute.Tensor, - g_out: cute.Tensor, - stream: CUstream, - ) -> None: - hidden_size = cute.size(g_out, mode=[1]) - copy_lora = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - self.element_type, - num_bits_per_copy=_VECTOR_WIDTH * self.element_type.width, - load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, - ) - copy_weight = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - self.element_type, - num_bits_per_copy=_VECTOR_WIDTH * self.element_type.width, - load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, - ) - self.kernel( - g_lora, - g_weight, - g_x, - g_out, - hidden_size, - copy_lora, - copy_weight, - ).launch( - grid=[cute.ceil_div(hidden_size, _OUTPUTS_PER_BLOCK), 1, 1], - block=[_BLOCK_SIZE, 1, 1], - stream=stream, - min_blocks_per_mp=1, - ) - - @cute.kernel - def kernel( - self, - g_lora: cute.Tensor, - g_weight: cute.Tensor, - g_x: cute.Tensor, - g_out: cute.Tensor, - hidden_size: cutlass.Int32, - copy_lora: cute.CopyAtom, - copy_weight: cute.CopyAtom, - ) -> None: - tidx, _, _ = cute.arch.thread_idx() - block_idx, _, _ = cute.arch.block_idx() - - acc = cute.make_rmem_tensor( - cute.make_layout( - (_OUTPUTS_PER_BLOCK, _HC), - stride=(_HC, 1), - ), - cutlass.Float32, - ) - acc.fill(0.0) - - lora_vec = cute.logical_divide(g_lora, (None, _VECTOR_WIDTH)) - weight_vec = cute.logical_divide(g_weight, (None, _VECTOR_WIDTH)) - lora_tiles = cute.logical_divide(lora_vec, (None, (None, _BLOCK_SIZE))) - weight_tiles = cute.logical_divide(weight_vec, (None, (None, _BLOCK_SIZE))) - thread_lora = lora_tiles[None, (None, (tidx, None))] - lora_regs = cute.make_rmem_tensor( - cute.make_layout((_VECTOR_WIDTH,), stride=(1,)), - self.element_type, - ) - weight_regs = cute.make_rmem_tensor( - cute.make_layout( - (_OUTPUTS_PER_BLOCK, _HC, _VECTOR_WIDTH), - stride=(_HC * _VECTOR_WIDTH, _VECTOR_WIDTH, 1), - ), - self.element_type, - ) - - hidden_base = block_idx * _OUTPUTS_PER_BLOCK - for k_tile in cutlass.range_constexpr(_K // (_BLOCK_SIZE * _VECTOR_WIDTH)): - cute.copy(copy_lora, thread_lora[0, None, k_tile], lora_regs) - for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): - for hc_stream in cutlass.range_constexpr(_HC): - weight_row = hc_stream * hidden_size + hidden_base + output - thread_weight = weight_tiles[weight_row, (None, (tidx, None))] - cute.copy( - copy_weight, - thread_weight[None, k_tile], - weight_regs[output, hc_stream, None], - ) - - raw_lora = lora_regs.load().to(cutlass.Float32) / _HC - activated_lora = (raw_lora / (1.0 + cute.exp(-raw_lora, fastmath=True))).to( - self.element_type - ) - activated_lora = activated_lora.to(cutlass.Float32) - weights = weight_regs.load().to(cutlass.Float32) - for vector_lane in cutlass.range_constexpr(_VECTOR_WIDTH): - for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): - for hc_stream in cutlass.range_constexpr(_HC): - acc[output, hc_stream] += ( - activated_lora[vector_lane] - * weights[output, hc_stream, vector_lane] - ) - - for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): - for hc_stream in cutlass.range_constexpr(_HC): - acc[output, hc_stream] = cute.arch.warp_reduction_sum( - acc[output, hc_stream] - ) - - if tidx == 0: - for output in cutlass.range_constexpr(_OUTPUTS_PER_BLOCK): - hidden_idx = hidden_base + output - if hidden_idx < hidden_size: - mixed = cutlass.Float32(0.0) - for hc_stream in cutlass.range_constexpr(_HC): - x = g_x[0, hc_stream * hidden_size + hidden_idx].to( - cutlass.Float32 - ) - mixed += x / ( - 1.0 + cute.exp(-acc[output, hc_stream], fastmath=True) - ) - g_out[0, hidden_idx] = (mixed / _HC).to(self.element_type) diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py b/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py deleted file mode 100644 index bcd87e938af3..000000000000 --- a/vllm/models/qwen4_exp/nvidia/ops/hc_silu_up_gate_mix.py +++ /dev/null @@ -1,185 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from typing import Any, ClassVar - -import torch - -from vllm.model_executor.warmup.cutedsl_warmup import ( - CuTeDSLCompileUnit, - register_cutedsl_warmup_provider, -) -from vllm.platforms import current_platform -from vllm.utils.torch_utils import current_stream, direct_register_custom_op - -_HC = 4 -_HIDDEN_SIZE = 2560 -_HYPER_HIDDEN_SIZE = _HC * _HIDDEN_SIZE -_LORA_RANK = 320 - - -class HCSiluUpGateMixOp: - """Process-local compiled fused HyperConnection decode operation.""" - - _instance: ClassVar[HCSiluUpGateMixOp | None] = None - - @classmethod - def initialize(cls) -> HCSiluUpGateMixOp: - if cls._instance is None: - cls._instance = cls() - return cls._instance - - @staticmethod - def is_supported(dtype: torch.dtype) -> bool: - if dtype != torch.bfloat16 or not current_platform.is_device_capability((9, 0)): - return False - try: - import cutlass # noqa: F401 - import cutlass.cute # noqa: F401 - except ImportError: - return False - return True - - def __init__(self) -> None: - self._compiled: Any | None = None - register_cutedsl_warmup_provider(self) - - def _compile(self) -> None: - import cutlass - import cutlass.cute as cute - from cuda.bindings.driver import CUstream - from quack.compile_utils import make_fake_tensor - - from ._hc_silu_up_gate_mix import HCSiluUpGateMixKernel - - lora = make_fake_tensor( - cutlass.BFloat16, - (1, _LORA_RANK), - divisibility=2, - ) - weight = make_fake_tensor( - cutlass.BFloat16, - (_HYPER_HIDDEN_SIZE, _LORA_RANK), - divisibility=2, - ) - x = make_fake_tensor( - cutlass.BFloat16, - (1, _HYPER_HIDDEN_SIZE), - divisibility=1, - ) - out = make_fake_tensor( - cutlass.BFloat16, - (1, _HIDDEN_SIZE), - divisibility=1, - ) - self._compiled = cute.compile( - HCSiluUpGateMixKernel(), - lora, - weight, - x, - out, - CUstream(current_stream().cuda_stream), - options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", - ) - - def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: - return ( - CuTeDSLCompileUnit( - name="Qwen4Exp fused HC SiLU/up-projection/gate-mix", - key=("qwen4-exp-hc-silu-up-gate-mix", torch.bfloat16), - compile=self._compile, - ), - ) - - def __call__( - self, - lora: torch.Tensor, - weight: torch.Tensor, - x: torch.Tensor, - ) -> torch.Tensor: - self._validate_inputs(lora, weight, x) - if self._compiled is None: - self._compile() - compiled = self._compiled - assert compiled is not None - out = torch.empty( - (1, _HIDDEN_SIZE), - dtype=torch.bfloat16, - device=lora.device, - ) - from cuda.bindings.driver import CUstream - - compiled( - lora, - weight, - x, - out, - CUstream(current_stream().cuda_stream), - ) - return out - - @staticmethod - def _validate_inputs( - lora: torch.Tensor, - weight: torch.Tensor, - x: torch.Tensor, - ) -> None: - if lora.shape != (1, _LORA_RANK): - raise ValueError(f"lora must have shape [1, {_LORA_RANK}].") - if weight.shape != (_HYPER_HIDDEN_SIZE, _LORA_RANK): - raise ValueError( - f"weight must have shape [{_HYPER_HIDDEN_SIZE}, {_LORA_RANK}]." - ) - if x.shape != (1, _HYPER_HIDDEN_SIZE): - raise ValueError(f"x must have shape [1, {_HYPER_HIDDEN_SIZE}].") - - tensors = (lora, weight, x) - if any(tensor.dtype != torch.bfloat16 for tensor in tensors): - raise ValueError("All inputs must use torch.bfloat16.") - if any(tensor.device != lora.device for tensor in tensors): - raise ValueError("All inputs must be on the same device.") - if any(not tensor.is_contiguous() for tensor in tensors): - raise ValueError("All inputs must be contiguous.") - - -def _hc_silu_up_gate_mix( - lora: torch.Tensor, - weight: torch.Tensor, - x: torch.Tensor, -) -> torch.Tensor: - if lora.shape[0] == 1: - return HCSiluUpGateMixOp.initialize()(lora, weight, x) - - from .hc import hc_gate_mix, hc_silu - - gate = torch.nn.functional.linear(hc_silu(lora, _HC), weight) - return hc_gate_mix(x, gate, _HC) - - -def _hc_silu_up_gate_mix_fake( - lora: torch.Tensor, - weight: torch.Tensor, - x: torch.Tensor, -) -> torch.Tensor: - del weight, x - return lora.new_empty((lora.shape[0], _HIDDEN_SIZE)) - - -direct_register_custom_op( - op_name="qwen4_exp_hc_silu_up_gate_mix", - op_func=_hc_silu_up_gate_mix, - fake_impl=_hc_silu_up_gate_mix_fake, -) - - -def hc_silu_up_gate_mix( - lora: torch.Tensor, - weight: torch.Tensor, - x: torch.Tensor, -) -> torch.Tensor: - return torch.ops.vllm.qwen4_exp_hc_silu_up_gate_mix(lora, weight, x) - - -__all__ = ["HCSiluUpGateMixOp", "hc_silu_up_gate_mix"] From b447a4a77dd41071c3f2c614031bd19d49190e13 Mon Sep 17 00:00:00 2001 From: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:06:47 +0800 Subject: [PATCH 4/4] [Kernel][Qwen] Retune Hopper LL-GEMM plans Retune all Qwen4Exp TP=4 skinny-GEMM shapes on H200 for M=1,2,4,8,16. Keep only plans that beat the standard linear path in both hot-cache and L2-flush measurements, and fall back for the remaining points. Add table integrity checks and H200 correctness coverage for every selected plan. Assisted-by: OpenAI Codex Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com> --- tests/kernels/test_bf16_skinny_gemm.py | 45 +++++++++++-- .../qwen4_exp/nvidia/low_latency_gemm.py | 65 ++++++++++++++++++- 2 files changed, 103 insertions(+), 7 deletions(-) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 9021b1d640b0..ca4e8c08e453 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -78,6 +78,12 @@ for _, config in spec.cute_configs ] +QWEN4_EXP_SM90_CASES = [ + (n, k, num_tokens, config) + for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS.items() + for num_tokens, config in plans.items() +] + EXPECTED_CUTE_CONFIGS = { (3072, 7168, 1): (224, 3, 4, 8), (3072, 7168, 2): (128, 3, 2, 8), @@ -529,11 +535,19 @@ def test_low_latency_table_capability_routing( assert k3_gemm._low_latency_table() is None -def test_qwen4_exp_hopper_plans_are_decode_only() -> None: +def test_qwen4_exp_hopper_plans_are_valid() -> None: plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS - assert plans.keys() == qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS.keys() - assert all(set(shape_plans) == {1} for shape_plans in plans.values()) + assert len(plans) == 9 + assert sum(map(len, plans.values())) == 31 + assert (320, 10240) in plans + assert (10240, 320) not in plans + for (n, k), shape_plans in plans.items(): + for num_tokens, config in shape_plans.items(): + assert config.num_rows == num_tokens + assert n % config.outputs_per_block == 0 + assert k % (config.block_size * config.vector_width) == 0 + assert config.static_k in (None, k) @pytest.mark.parametrize( @@ -673,7 +687,7 @@ def _require_capability_and_cute(capability: tuple[int, int]) -> None: not torch.cuda.is_available() or torch.cuda.get_device_capability() != capability ): - pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}") + pytest.skip(f"CuTe DSL selection requires SM{capability[0]}{capability[1]}") if not k3_gemm.shape_dynamic_skinny_gemm.is_available(): pytest.skip("CuTe DSL is not available") @@ -709,6 +723,29 @@ def test_glm_cute_selected_shapes( torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) +@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM90_CASES) +def test_qwen4_exp_sm90_selected_shapes( + n: int, + k: int, + num_tokens: int, + config: SkinnyGemmConfig, +) -> None: + _require_capability_and_cute((9, 0)) + torch.manual_seed(42 + num_tokens) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + selected = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS[(n, k)][num_tokens] + assert selected == config + output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight) + + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + def test_glm52_q_b_nonpacked_single_row_falls_back() -> None: _require_sm103_and_cute() spec = glm52_gemm.GLM52_Q_B_PROJECTION diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index 4f963f6bc082..c177f077451e 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -78,10 +78,69 @@ }, } -# H200 plans measured under CUDA graph replay. Only M=1 is enabled so larger -# batches retain the standard linear implementation and its GEMM heuristics. +# H200 plans selected by exhaustive CUDA graph replay measurements over +# M={1, 2, 4, 8, 16}. Only points that beat the standard linear implementation +# in both hot-cache and L2-flush measurements are retained; other token counts +# keep the standard implementation and its GEMM heuristics. QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { - shape: {1: plans[1]} for shape, plans in QWEN4_EXP_GEMM_PLANS.items() + # GDN fused QKVZ projection, TP=4. + (4096, 2560): { + 1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560), + 2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=2560), + }, + # GDN and QSA output projections, TP=4. + (2560, 1536): { + 1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=1536), + 2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=1536), + 4: SkinnyGemmConfig(4, 64, 4, k_unroll=6, vector_width=4), + }, + # GDN fused B/A projection, TP=4. + (24, 2560): { + 1: SkinnyGemmConfig(1, 128, 3, vector_width=4, static_k=2560), + 2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=2560), + 4: SkinnyGemmConfig(4, 64, 1, static_k=2560), + 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), + 2: SkinnyGemmConfig(2, 128, 4, k_unroll=5, vector_width=4), + 4: SkinnyGemmConfig(4, 160, 1, k_unroll=2), + 8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560), + 16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560), + }, + # LM head, TP=4. + (62080, 2560): { + 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), + 2: SkinnyGemmConfig(2, 128, 1, k_unroll=10), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=10), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=10), + }, }