[WS2][gemm][2th merge] ffn collectives on test

csrc/cuda/gemm/det_gemm_kernel.cu

@@ -78,6 +78,42 @@ __device__ nv_bf16 k_tree_naive(const nv_bf16* __restrict__ A, const nv_bf16* __
                   k_tree_naive(A, B, row, col, N, K, mid, hi));
 }
 
+// Must match SM90 BK so an aligned-K naive tree equals the SM90 tile tree.
+constexpr int K_TREE_LEAF = 32;
+
+__device__ __forceinline__ nv_bf16 bf16_add(nv_bf16 a, nv_bf16 b) {
+  return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b));
+}
+
+// True iff [lo, hi) is a node of the mid-split tree over [0, n).
+__device__ __forceinline__ bool is_mid_split_node(int lo, int hi, int n) {
+  int a = 0, b = n;
+  while (b - a > 1) {
+    if (a == lo && b == hi) return true;
+    const int m = a + (b - a) / 2;
+    if (hi <= m)
+      b = m;
+    else if (lo >= m)
+      a = m;
+    else
+      return false;
+  }
+  return a == lo && b == hi;
+}
+
+__device__ nv_bf16 k_tree_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B,
+                                int row, int col, int N, int K, int lo, int hi) {
+  if (hi - lo <= K_TREE_LEAF) {
+    float acc = 0.0f;
+    for (int k = lo; k < hi; ++k)
+      acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]);
+    return __float2bfloat16(acc);
+  }
+  const int mid = lo + (hi - lo) / 2;
+  return bf16_add(k_tree_naive(A, B, row, col, N, K, lo, mid),
+                  k_tree_naive(A, B, row, col, N, K, mid, hi));
+}
+
 // Naive FP32 scalar kernel (fallback + ground truth). Batch-invariant by
 // construction: one thread = one output element, mid-split K tree.
 constexpr int NAIVE_TILE = 16;

rl_engine/kernels/ops/pytorch/ffn/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026 RL-Kernel Contributors
 
-from .ffn import BACKEND_ID, Qwen3FFNOp, qwen3_ffn
+from .ffn import qwen3_ffn
 
-__all__ = ["BACKEND_ID", "Qwen3FFNOp", "qwen3_ffn"]
+__all__ = ["qwen3_ffn"]

rl_engine/kernels/ops/pytorch/ffn/ffn.py

@@ -13,47 +13,25 @@
 
 QWEN3_8B_HIDDEN_SIZE = 4096
 QWEN3_8B_INTERMEDIATE_SIZE = 12288
-BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1"
 
-_DET_GEMM_SYMBOLS = (
+_REQUIRED_SYMBOLS = (
     "det_gemm_fwd",
     "det_gemm_db",
-)
-_SWIGLU_SYMBOLS = (
     "swiglu_forward",
     "swiglu_backward",
 )
-_REQUIRED_SYMBOLS = _DET_GEMM_SYMBOLS + _SWIGLU_SYMBOLS
 _COLLECTIVE_MIN_CAPACITY_BYTES = 64 * 1024 * 1024
 _COLLECTIVES: dict[tuple[int, int, int, int], Any] = {}
 
 
-def _require_ffn_kernels(*, disable_split_k: bool) -> None:
-    required = _REQUIRED_SYMBOLS if disable_split_k else _SWIGLU_SYMBOLS
-    missing = [name for name in required if not hasattr(_C, name)]
+def _require_ffn_kernels() -> None:
+    missing = [name for name in _REQUIRED_SYMBOLS if not hasattr(_C, name)]
     if not _EXT_AVAILABLE or _C is None or missing:
         suffix = f" Missing symbols: {', '.join(missing)}." if missing else ""
-        needed = (
-            "compiled deterministic GEMM and SwiGLU CUDA kernels"
-            if disable_split_k
-            else "compiled SwiGLU CUDA kernels"
+        raise RuntimeError(
+            "qwen3_ffn requires the compiled deterministic GEMM and "
+            f"SwiGLU CUDA kernels.{suffix}"
         )
-        raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}")
-
-
-def _gemm_fwd(a: Tensor, b: Tensor, *, disable_split_k: bool) -> Tensor:
-    if disable_split_k:
-        return _C.det_gemm_fwd(a, b)
-    # cuBLASLt / CUTLASS: may use split-K. Detach so Autograd.Function owns backward.
-    with torch.no_grad():
-        return torch.matmul(a, b)
-
-
-def _gemm_db(a: Tensor, grad_output: Tensor, *, disable_split_k: bool) -> Tensor:
-    if disable_split_k:
-        return _C.det_gemm_db(a, grad_output)
-    with torch.no_grad():
-        return torch.matmul(a.t().contiguous(), grad_output)
 
 
 def _require_parallel_group(group: Any, name: str):
@@ -127,10 +105,10 @@ def _collective_for_group(group: Any, *, min_size_bytes: int):
     if group is None:
         return None
 
-    import torch.distributed as dist
-
     from rl_engine.distributed import DeterministicCollective
 
+    import torch.distributed as dist
+
     rank = dist.get_rank(group=group)
     world_size = dist.get_world_size(group=group)
     device_index = torch.cuda.current_device()
@@ -178,7 +156,6 @@ def forward(
         tp_group: Any,
         cp_group: Any,
         sequence_parallel: bool,
-        disable_split_k: bool,
     ) -> Tensor:
         tp_dist = _require_parallel_group(tp_group, "tensor")
         _require_parallel_group(cp_group, "context")
@@ -205,18 +182,10 @@ def forward(
             rmsnorm_output_2d = _all_gather_tokens(rmsnorm_output_2d, tp_collective)
 
         # The model stores projection weights as [out, in]; GEMM consumes [K, N].
-        gate = _gemm_fwd(
-            rmsnorm_output_2d,
-            gate_weight.t().contiguous(),
-            disable_split_k=disable_split_k,
-        )
-        up = _gemm_fwd(
-            rmsnorm_output_2d,
-            up_weight.t().contiguous(),
-            disable_split_k=disable_split_k,
-        )
+        gate = _C.det_gemm_fwd(rmsnorm_output_2d, gate_weight.t().contiguous())
+        up = _C.det_gemm_fwd(rmsnorm_output_2d, up_weight.t().contiguous())
         activated = _C.swiglu_forward(gate, up)
-        output = _gemm_fwd(activated, down_weight.t().contiguous(), disable_split_k=disable_split_k)
+        output = _C.det_gemm_fwd(activated, down_weight.t().contiguous())
 
         if sequence_parallel:
             output = _reduce_scatter_tokens(output, tp_collective)
@@ -236,7 +205,6 @@ def forward(
         ctx.tp_collective = tp_collective
         ctx.cp_collective = cp_collective
         ctx.sequence_parallel = sequence_parallel
-        ctx.disable_split_k = disable_split_k
         return output.reshape(*input_shape[:-1], output.size(-1))
 
     @staticmethod
@@ -252,7 +220,6 @@ def backward(ctx, grad_output: Tensor):
         ) = ctx.saved_tensors
         tp_collective = ctx.tp_collective
         cp_collective = ctx.cp_collective
-        disable_split_k = ctx.disable_split_k
         grad_output = grad_output.reshape(-1, grad_output.size(-1)).contiguous()
         if ctx.sequence_parallel:
             grad_output = _all_gather_tokens(grad_output, tp_collective)
@@ -263,46 +230,26 @@ def backward(ctx, grad_output: Tensor):
         if cp_collective is not None:
             activated_full = _all_gather_tokens(activated, cp_collective)
             grad_output_full = _all_gather_tokens(grad_output, cp_collective)
-            grad_down_weight = (
-                _gemm_db(activated_full, grad_output_full, disable_split_k=disable_split_k)
-                .t()
-                .contiguous()
-            )
+            grad_down_weight = _C.det_gemm_db(activated_full, grad_output_full).t().contiguous()
         else:
-            grad_down_weight = (
-                _gemm_db(activated, grad_output, disable_split_k=disable_split_k).t().contiguous()
-            )
+            grad_down_weight = _C.det_gemm_db(activated, grad_output).t().contiguous()
 
         # Down input-gradient shards concatenate across TP; no TP reduction.
-        grad_activated = _gemm_fwd(grad_output, down_weight, disable_split_k=disable_split_k)
+        grad_activated = _C.det_gemm_fwd(grad_output, down_weight)
         grad_gate, grad_up = _C.swiglu_backward(grad_activated, gate, up)
 
         if cp_collective is not None:
             rmsnorm_full = _all_gather_tokens(rmsnorm_output, cp_collective)
             grad_gate_full = _all_gather_tokens(grad_gate, cp_collective)
             grad_up_full = _all_gather_tokens(grad_up, cp_collective)
-            grad_gate_weight = (
-                _gemm_db(rmsnorm_full, grad_gate_full, disable_split_k=disable_split_k)
-                .t()
-                .contiguous()
-            )
-            grad_up_weight = (
-                _gemm_db(rmsnorm_full, grad_up_full, disable_split_k=disable_split_k)
-                .t()
-                .contiguous()
-            )
+            grad_gate_weight = _C.det_gemm_db(rmsnorm_full, grad_gate_full).t().contiguous()
+            grad_up_weight = _C.det_gemm_db(rmsnorm_full, grad_up_full).t().contiguous()
         else:
-            grad_gate_weight = (
-                _gemm_db(rmsnorm_output, grad_gate, disable_split_k=disable_split_k)
-                .t()
-                .contiguous()
-            )
-            grad_up_weight = (
-                _gemm_db(rmsnorm_output, grad_up, disable_split_k=disable_split_k).t().contiguous()
-            )
+            grad_gate_weight = _C.det_gemm_db(rmsnorm_output, grad_gate).t().contiguous()
+            grad_up_weight = _C.det_gemm_db(rmsnorm_output, grad_up).t().contiguous()
 
         # Gate/Up input gradients reduce across TP, then add locally.
-        grad_rmsnorm_from_gate = _gemm_fwd(grad_gate, gate_weight, disable_split_k=disable_split_k)
+        grad_rmsnorm_from_gate = _C.det_gemm_fwd(grad_gate, gate_weight)
         if ctx.sequence_parallel:
             grad_rmsnorm_from_gate = _reduce_scatter_tokens(
                 grad_rmsnorm_from_gate,
@@ -314,7 +261,7 @@ def backward(ctx, grad_output: Tensor):
                 tp_collective,
             )
 
-        grad_rmsnorm_from_up = _gemm_fwd(grad_up, up_weight, disable_split_k=disable_split_k)
+        grad_rmsnorm_from_up = _C.det_gemm_fwd(grad_up, up_weight)
         if ctx.sequence_parallel:
             grad_rmsnorm_from_up = _reduce_scatter_tokens(
                 grad_rmsnorm_from_up,
@@ -335,7 +282,6 @@ def backward(ctx, grad_output: Tensor):
             None,
             None,
             None,
-            None,
         )
 
 
@@ -348,8 +294,6 @@ def qwen3_ffn(
     tp_group: Any = None,
     cp_group: Any = None,
     sequence_parallel: bool = False,
-    deterministic: bool | None = None,
-    disable_split_k: bool | None = None,
 ) -> Tensor:
     """Apply a bias-free SiLU-gated FFN with deterministic backward kernels.
 
@@ -372,20 +316,14 @@ def qwen3_ffn(
             are sharded on the flattened token dimension across ``tp_group``.
             Token gather/scatter use the deterministic AllGather and
             ReduceScatter.
-        deterministic: Select the RL-Kernel fixed-reduction GEMM when True
-            (default), or the production ``torch.matmul`` GEMM when False.
-        disable_split_k: Compatibility alias for ``deterministic``. New code
-            should use ``deterministic`` because Split-K is only one possible
-            implementation detail of the production GEMM.
 
     Returns:
         FFN output with shape ``[..., H]``.
     """
+    _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight)
+    _require_ffn_kernels()
     if not isinstance(sequence_parallel, bool):
         raise TypeError("sequence_parallel must be a bool.")
-    deterministic = _resolve_deterministic_mode(deterministic, disable_split_k)
-    _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight)
-    _require_ffn_kernels(disable_split_k=deterministic)
     return _DeterministicFFNFunction.apply(
         rmsnorm_output,
         gate_weight,
@@ -394,84 +332,4 @@ def qwen3_ffn(
         tp_group,
         cp_group,
         sequence_parallel,
-        deterministic,
     )
-
-
-def _resolve_deterministic_mode(
-    deterministic: bool | None,
-    disable_split_k: bool | None,
-) -> bool:
-    if deterministic is not None and not isinstance(deterministic, bool):
-        raise TypeError("deterministic must be a bool or None.")
-    if disable_split_k is not None and not isinstance(disable_split_k, bool):
-        raise TypeError("disable_split_k must be a bool or None.")
-    if (
-        deterministic is not None
-        and disable_split_k is not None
-        and deterministic != disable_split_k
-    ):
-        raise ValueError("deterministic and disable_split_k select conflicting FFN backends.")
-    if deterministic is not None:
-        return deterministic
-    if disable_split_k is not None:
-        return disable_split_k
-    return True
-
-
-class Qwen3FFNOp:
-    """Instantiable Qwen3 FFN wrapper for semantic operator dispatch."""
-
-    op_class = "ffn"
-    is_batch_invariant = True
-    backend_id = BACKEND_ID
-
-    def __call__(
-        self,
-        rmsnorm_output: Tensor,
-        gate_weight: Tensor,
-        up_weight: Tensor,
-        down_weight: Tensor,
-        *,
-        tp_group: Any = None,
-        cp_group: Any = None,
-        sequence_parallel: bool = False,
-        deterministic: bool | None = None,
-        disable_split_k: bool | None = None,
-    ) -> Tensor:
-        return self.apply(
-            rmsnorm_output,
-            gate_weight,
-            up_weight,
-            down_weight,
-            tp_group=tp_group,
-            cp_group=cp_group,
-            sequence_parallel=sequence_parallel,
-            deterministic=deterministic,
-            disable_split_k=disable_split_k,
-        )
-
-    def apply(
-        self,
-        rmsnorm_output: Tensor,
-        gate_weight: Tensor,
-        up_weight: Tensor,
-        down_weight: Tensor,
-        *,
-        tp_group: Any = None,
-        cp_group: Any = None,
-        sequence_parallel: bool = False,
-        deterministic: bool | None = None,
-        disable_split_k: bool | None = None,
-    ) -> Tensor:
-        return qwen3_ffn(
-            rmsnorm_output,
-            gate_weight,
-            up_weight,
-            down_weight,
-            tp_group=tp_group,
-            cp_group=cp_group,
-            sequence_parallel=sequence_parallel,
-            deterministic=deterministic,
-            disable_split_k=disable_split_k,
-        )