perf(qwen4): optimize gated residual kernels and PDL

python/tokenspeed/runtime/layers/hyperconnection.py

@@ -26,9 +26,12 @@
 
 import torch
 import torch.nn.functional as F
-import triton
-import triton.language as tl
-from tokenspeed_kernel.platform import pdl_enabled
+from tokenspeed_kernel import (
+    gated_residual_combine,
+    gated_residual_mix,
+    grouped_gemma_rmsnorm,
+    prepare_gated_residual_weight_cache,
+)
 from torch import nn
 
 
@@ -43,6 +46,12 @@ class HyperConnectionConfig:
     params_dtype: torch.dtype = torch.bfloat16
     hc_per_branch_norm: bool = True
 
+    def __post_init__(self) -> None:
+        if self.hc_count <= 1:
+            raise ValueError("hc_count must be greater than one")
+        if self.hidden_size <= 0 or self.hc_lowrank <= 0:
+            raise ValueError("hidden_size and hc_lowrank must be positive")
+
 
 def _matching_rows(base: torch.Tensor, derived: torch.Tensor):
     """Row offset of ``derived`` inside ``base``, or ``None`` when unrelated.
@@ -70,7 +79,7 @@ def _matching_rows(base: torch.Tensor, derived: torch.Tensor):
 
 
 class GroupedGemmaRMSNorm(nn.Module):
-    """Gemma RMSNorm with optional independently-normalized feature groups."""
+    """GPU Gemma RMSNorm with optional independently-normalized feature groups."""
 
     def __init__(self, hidden_size: int, eps: float, group_size: int | None = None):
         super().__init__()
@@ -90,123 +99,13 @@ def _weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> No
         param.data.copy_(loaded_weight)
         self.gemma_weight = param.data + 1.0
 
-    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
-        input_dtype = x.dtype
-        x_float = x.float()
-        if self.group_size is None:
-            variance = x_float.square().mean(dim=-1, keepdim=True)
-            normalized = x_float * torch.rsqrt(variance + self.variance_epsilon)
-        else:
-            grouped = x_float.unflatten(-1, (-1, self.group_size))
-            variance = grouped.square().mean(dim=-1, keepdim=True)
-            normalized = (
-                grouped * torch.rsqrt(variance + self.variance_epsilon)
-            ).flatten(-2)
-        return (normalized * self.gemma_weight.float()).to(input_dtype)
-
-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        from tokenspeed.runtime.layers.attention.linear.layernorm_gated import (
-            rmsnorm_fn,
-        )
-
-        return rmsnorm_fn(
-            x,
-            self.gemma_weight,
-            None,
-            z=None,
-            eps=self.variance_epsilon,
-            group_size=self.group_size,
-        )
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if x.is_cuda:
-            return self.forward_cuda(x)
-        return self.forward_native(x)
-
-
-@triton.jit
-def _hc_mix_epilogue_kernel(
-    gate_ptr,
-    normalized_ptr,
-    out_ptr,
-    gate_row_stride,
-    normalized_row_stride,
-    out_row_stride,
-    hidden_size,
-    HC_COUNT: tl.constexpr,
-    BLOCK: tl.constexpr,
-    ENABLE_PDL: tl.constexpr,
-):
-    """Sigmoid gate, per-branch weighting and branch mean in one pass.
-
-    Each program owns one token and one slice of ``hidden_size``, so the widened
-    ``[..., hc_count * hidden_size]`` sigmoid and product tensors the eager path
-    materializes stay in registers and only the mixed row reaches memory.
-    """
-
-    row = tl.program_id(0)
-    offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
-    mask = offsets < hidden_size
-    if ENABLE_PDL:
-        tl.extra.cuda.gdc_wait()
-    mixed = tl.zeros([BLOCK], dtype=tl.float32)
-    for branch in tl.static_range(HC_COUNT):
-        column = branch * hidden_size + offsets
-        gate = tl.load(
-            gate_ptr + row * gate_row_stride + column, mask=mask, other=0.0
-        ).to(tl.float32)
-        value = tl.load(
-            normalized_ptr + row * normalized_row_stride + column, mask=mask, other=0.0
-        ).to(tl.float32)
-        mixed += tl.sigmoid(gate) * value
-    tl.store(out_ptr + row * out_row_stride + offsets, mixed / HC_COUNT, mask=mask)
-    if ENABLE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
-
-
-@triton.jit
-def _hc_combine_kernel(
-    hyper_ptr,
-    block_ptr,
-    inject_ptr,
-    out_ptr,
-    hyper_row_stride,
-    block_row_stride,
-    inject_row_stride,
-    out_row_stride,
-    hidden_size,
-    HC_COUNT: tl.constexpr,
-    BLOCK: tl.constexpr,
-    ENABLE_PDL: tl.constexpr,
-):
-    """Gate one sublayer output and add it into every residual branch.
-
-    The sublayer row is loaded once and reused for all ``hc_count`` branches, so
-    the broadcast product the eager path materializes at the full residual width
-    never reaches memory.
-    """
-
-    row = tl.program_id(0)
-    offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
-    mask = offsets < hidden_size
-    if ENABLE_PDL:
-        tl.extra.cuda.gdc_wait()
-    value = tl.load(
-        block_ptr + row * block_row_stride + offsets, mask=mask, other=0.0
-    ).to(tl.float32)
-    for branch in tl.static_range(HC_COUNT):
-        logit = tl.load(inject_ptr + row * inject_row_stride + branch).to(tl.float32)
-        column = branch * hidden_size + offsets
-        residual = tl.load(
-            hyper_ptr + row * hyper_row_stride + column, mask=mask, other=0.0
-        ).to(tl.float32)
-        tl.store(
-            out_ptr + row * out_row_stride + column,
-            residual + value * 2.0 * tl.sigmoid(logit),
-            mask=mask,
+        return grouped_gemma_rmsnorm(
+            x,
+            self.weight,
+            self.group_size,
+            self.variance_epsilon,
         )
-    if ENABLE_PDL:
-        tl.extra.cuda.gdc_launch_dependents()
 
 
 class GatedResidualSimple(nn.Module):
@@ -237,8 +136,13 @@ def __init__(
         # once: rows [0, hc_lowrank) hold the mix-down weight and the next
         # hc_count rows the inject weight. Checkpoint shards are routed there by
         # the stacked mapping in load_qwen4_exp_weights. Both projections are
-        # scaled by 1 / hc_count, folded into the weight at load time, which is
-        # exact for a power-of-two hc_count.
+        # scaled by 1 / hc_count. Fold that scale only when it is an exact
+        # binary exponent change; non-power-of-two counts scale projection
+        # results at runtime instead of quantizing the checkpoint weights.
+        self._fold_projection_scale = not (self.hc_count & (self.hc_count - 1))
+        self._projection_scale = (
+            1.0 if self._fold_projection_scale else 1.0 / self.hc_count
+        )
         projection_rows = (config.hc_lowrank if use_mix else 0) + (
             self.hc_count if use_combine else 0
         )
@@ -250,6 +154,18 @@ def __init__(
         self.input_mix_weight_up = (
             nn.Linear(config.hc_lowrank, hc_size, bias=False) if use_mix else None
         )
+        if self.input_mix_weight_up is not None:
+            self.input_mix_weight_up.weight.weight_loader = self._load_up_weight
+
+    def _load_up_weight(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        """Load the fixed-shape projection and prepare its derived GPU weight."""
+        if param.shape != loaded_weight.shape:
+            raise ValueError(
+                f"hyper-connection up weight shape mismatch: param "
+                f"{tuple(param.shape)}, loaded {tuple(loaded_weight.shape)}"
+            )
+        param.data.copy_(loaded_weight)
+        prepare_gated_residual_weight_cache(param, self.hc_lowrank)
 
     def _load_projection_shard(
         self,
@@ -270,17 +186,12 @@ def _load_projection_shard(
                 f"hyper-connection {shard_id} shard shape mismatch: param "
                 f"{tuple(param.shape)}, loaded {tuple(loaded_weight.shape)}"
             )
-        param.data[start : start + size].copy_(
-            (loaded_weight / self.hc_count).to(param.device, param.dtype)
+        weight = (
+            loaded_weight / self.hc_count
+            if self._fold_projection_scale
+            else loaded_weight
         )
-
-    def _split_projection(self, projected: torch.Tensor):
-        """Split the fused projection into mix gate and inject logits."""
-        if not self.use_combine:
-            return projected, None
-        if not self.use_mix:
-            return None, projected
-        return projected[..., : self.hc_lowrank], projected[..., self.hc_lowrank :]
+        param.data[start : start + size].copy_(weight.to(param.device, param.dtype))
 
     def _inject_logits(self, normalized: torch.Tensor):
         """Inject logits alone, for residual rows ``mix`` never saw."""
@@ -289,7 +200,10 @@ def _inject_logits(self, normalized: torch.Tensor):
         weight = self.mix_inject_proj.weight
         if self.use_mix:
             weight = weight[self.hc_lowrank :]
-        return F.linear(normalized, weight)
+        logits = F.linear(normalized, weight)
+        if self._projection_scale != 1.0:
+            logits = logits * self._projection_scale
+        return logits
 
     def _normalize(self, hyper_input: torch.Tensor) -> torch.Tensor:
         if self.config.hc_per_branch_norm:
@@ -315,45 +229,6 @@ def norm_for(self, value: torch.Tensor, residuals):
         rows = slice(start, start + value.shape[0])
         return value, normalized[rows], inject_logits[rows]
 
-    def _mix_epilogue_torch(
-        self, gate: torch.Tensor, normalized: torch.Tensor
-    ) -> torch.Tensor:
-        weights = torch.sigmoid(gate).unflatten(-1, (self.hc_count, self.hidden_size))
-        branches = normalized.unflatten(-1, (self.hc_count, self.hidden_size))
-        return (weights * branches).mean(dim=-2).to(self.config.params_dtype)
-
-    def _mix_epilogue_cuda(
-        self, gate: torch.Tensor, normalized: torch.Tensor
-    ) -> torch.Tensor:
-        rows = gate.shape[0]
-        mixed = torch.empty(
-            (rows, self.hidden_size),
-            dtype=self.config.params_dtype,
-            device=gate.device,
-        )
-        # Row strides let the fused projection's split views feed the kernel
-        # without a contiguity copy; only the last dim must be dense.
-        if gate.stride(-1) != 1:
-            gate = gate.contiguous()
-        if normalized.stride(-1) != 1:
-            normalized = normalized.contiguous()
-        block = min(triton.next_power_of_2(self.hidden_size), 1024)
-        use_pdl = pdl_enabled()
-        _hc_mix_epilogue_kernel[(rows, triton.cdiv(self.hidden_size, block))](
-            gate,
-            normalized,
-            mixed,
-            gate.stride(0),
-            normalized.stride(0),
-            mixed.stride(0),
-            self.hidden_size,
-            HC_COUNT=self.hc_count,
-            BLOCK=block,
-            ENABLE_PDL=use_pdl,
-            **({"launch_pdl": True} if use_pdl else {}),
-        )
-        return mixed
-
     def mix(self, hyper_input: torch.Tensor):
         """Mix ``hc_count`` residual branches into one sublayer input.
 
@@ -367,84 +242,32 @@ def mix(self, hyper_input: torch.Tensor):
                 f"hyper input width must be {expected}, got {hyper_input.shape[-1]}"
             )
         normalized = self._normalize(hyper_input)
-        if hyper_input.shape[0] == 0:
-            mixed = hyper_input.new_empty((*hyper_input.shape[:-1], self.hidden_size))
-            inject_logits = hyper_input.new_empty(
-                (*hyper_input.shape[:-1], self.hc_count)
-            )
-            return mixed, (hyper_input, normalized, inject_logits)
-        gate, inject_logits = self._split_projection(self.mix_inject_proj(normalized))
-        gate = self.input_mix_weight_up(F.silu(gate))
-        if gate.is_cuda and gate.dim() == 2:
-            mixed = self._mix_epilogue_cuda(gate, normalized)
-        else:
-            mixed = self._mix_epilogue_torch(gate, normalized)
-        return mixed, (
-            hyper_input,
+        mixed, inject_logits = gated_residual_mix(
             normalized,
-            inject_logits,
-        )
-
-    def _combine_torch(
-        self,
-        block_output: torch.Tensor,
-        hyper_input: torch.Tensor,
-        inject_logits: torch.Tensor,
-    ) -> torch.Tensor:
-        inject = 2 * torch.sigmoid(inject_logits)
-        combined = hyper_input.unflatten(
-            -1, (self.hc_count, self.hidden_size)
-        ) + block_output.unsqueeze(-2) * inject.unsqueeze(-1)
-        return combined.flatten(-2).to(self.config.params_dtype)
-
-    def _combine_cuda(
-        self,
-        block_output: torch.Tensor,
-        hyper_input: torch.Tensor,
-        inject_logits: torch.Tensor,
-    ) -> torch.Tensor:
-        rows = block_output.shape[0]
-        combined = torch.empty(
-            (rows, self.hc_count * self.hidden_size),
-            dtype=self.config.params_dtype,
-            device=block_output.device,
+            self.mix_inject_proj.weight,
+            self.input_mix_weight_up.weight,
+            self.hc_count,
+            self.hidden_size,
+            self.hc_lowrank,
+            projection_scale=self._projection_scale,
         )
-        # Row strides let sliced residuals and the fused projection's inject view
-        # feed the kernel without a contiguity copy; only the last dim must be
-        # dense.
-        if hyper_input.stride(-1) != 1:
-            hyper_input = hyper_input.contiguous()
-        if block_output.stride(-1) != 1:
-            block_output = block_output.contiguous()
-        if inject_logits.stride(-1) != 1:
-            inject_logits = inject_logits.contiguous()
-        block = min(triton.next_power_of_2(self.hidden_size), 1024)
-        use_pdl = pdl_enabled()
-        _hc_combine_kernel[(rows, triton.cdiv(self.hidden_size, block))](
+        mixed = mixed.to(self.config.params_dtype)
+        return mixed, (
             hyper_input,
-            block_output,
+            normalized,
             inject_logits,
-            combined,
-            hyper_input.stride(0),
-            block_output.stride(0),
-            inject_logits.stride(0),
-            combined.stride(0),
-            self.hidden_size,
-            HC_COUNT=self.hc_count,
-            BLOCK=block,
-            ENABLE_PDL=use_pdl,
-            **({"launch_pdl": True} if use_pdl else {}),
         )
-        return combined
 
     def combine(self, block_output: torch.Tensor, residuals) -> torch.Tensor:
         """Inject one sublayer output back into every residual branch."""
         hyper_input, _, inject_logits = residuals
-        if block_output.shape[0] == 0:
-            return hyper_input.to(self.config.params_dtype)
-        if block_output.is_cuda and block_output.dim() == 2:
-            return self._combine_cuda(block_output, hyper_input, inject_logits)
-        return self._combine_torch(block_output, hyper_input, inject_logits)
+        return gated_residual_combine(
+            block_output,
+            hyper_input,
+            inject_logits,
+            self.hc_count,
+            self.hidden_size,
+        ).to(self.config.params_dtype)
 
 
 __all__ = [