diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 174ea77c..ca494243 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/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; diff --git a/rl_engine/kernels/ops/pytorch/ffn/__init__.py b/rl_engine/kernels/ops/pytorch/ffn/__init__.py index 19a54d23..eb8a4871 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/__init__.py +++ b/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"] diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index c2507a44..05234bfc 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/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, - ) diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index 9b2b1d21..6cd32469 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -41,20 +41,29 @@ _WORLD2_CONFIGS = ( ("tp2_sp", 2, 1, True, _TOPOLOGY_TOKENS), ("cp2", 1, 2, False, _TOPOLOGY_TOKENS), - *((f"cp2_T{token_count}", 1, 2, False, token_count) for token_count in _CP_TOKEN_COUNTS), + *( + (f"cp2_T{token_count}", 1, 2, False, token_count) + for token_count in _CP_TOKEN_COUNTS + ), ) _WORLD4_CONFIGS = ( ("tp4", 4, 1, False, _TOPOLOGY_TOKENS), ("cp4", 1, 4, False, _TOPOLOGY_TOKENS), ("tp2_cp2", 2, 2, False, _TOPOLOGY_TOKENS), ("tp2_cp2_sp", 2, 2, True, _TOPOLOGY_TOKENS), - *((f"cp4_T{token_count}", 1, 4, False, token_count) for token_count in _CP_TOKEN_COUNTS), + *( + (f"cp4_T{token_count}", 1, 4, False, token_count) + for token_count in _CP_TOKEN_COUNTS + ), ) _WORLD8_WORLD_GROUP_CONFIGS = ( ("tp8", 8, 1, False, _TOPOLOGY_TOKENS), ("tp8_sp", 8, 1, True, _TOPOLOGY_TOKENS), ("cp8", 1, 8, False, _TOPOLOGY_TOKENS), - *((f"cp8_T{token_count}", 1, 8, False, token_count) for token_count in _CP_TOKEN_COUNTS), + *( + (f"cp8_T{token_count}", 1, 8, False, token_count) + for token_count in _CP_TOKEN_COUNTS + ), ) _WORLD8_TP2_CP4_CONFIGS = (("tp2_cp4", 2, 4, False, _TOPOLOGY_TOKENS),) _WORLD8_TP4_CP2_CONFIGS = ( @@ -157,7 +166,9 @@ def _shard_ranges( def _spawn_nccl_workers(worker, world_size: int, worker_args=(), *, timeout: int = 180) -> None: if not _has_sm90_ffn_devices(world_size): - pytest.skip(f"requires {world_size} SM90 GPUs, NCCL, and the GEMM/SwiGLU extension") + pytest.skip( + f"requires {world_size} SM90 GPUs, NCCL, and the GEMM/SwiGLU extension" + ) ctx = mp.get_context("spawn") with tempfile.TemporaryDirectory() as tmpdir: @@ -242,22 +253,13 @@ def _distributed_ffn_backward_nccl_worker( (token_count, hidden_size), seed=40, device=device, dtype=torch.bfloat16 ) gate_weight = _randn( - (intermediate_size, hidden_size), - seed=41, - device=device, - dtype=torch.bfloat16, + (intermediate_size, hidden_size), seed=41, device=device, dtype=torch.bfloat16 ) up_weight = _randn( - (intermediate_size, hidden_size), - seed=42, - device=device, - dtype=torch.bfloat16, + (intermediate_size, hidden_size), seed=42, device=device, dtype=torch.bfloat16 ) down_weight = _randn( - (hidden_size, intermediate_size), - seed=43, - device=device, - dtype=torch.bfloat16, + (hidden_size, intermediate_size), seed=43, device=device, dtype=torch.bfloat16 ) grad_output = _randn( (token_count, hidden_size), seed=44, device=device, dtype=torch.bfloat16 @@ -359,24 +361,17 @@ def _tp1_vs_tpn_train_infer_worker(rank, world_size, init_method, result_queue, ) device = torch.device("cuda", rank) token_count, hidden_size, intermediate_size = 16, 64, 256 - hidden = _randn((token_count, hidden_size), seed=50, device=device, dtype=torch.bfloat16) + hidden = _randn( + (token_count, hidden_size), seed=50, device=device, dtype=torch.bfloat16 + ) gate_weight = _randn( - (intermediate_size, hidden_size), - seed=51, - device=device, - dtype=torch.bfloat16, + (intermediate_size, hidden_size), seed=51, device=device, dtype=torch.bfloat16 ) up_weight = _randn( - (intermediate_size, hidden_size), - seed=52, - device=device, - dtype=torch.bfloat16, + (intermediate_size, hidden_size), seed=52, device=device, dtype=torch.bfloat16 ) down_weight = _randn( - (hidden_size, intermediate_size), - seed=53, - device=device, - dtype=torch.bfloat16, + (hidden_size, intermediate_size), seed=53, device=device, dtype=torch.bfloat16 ) grad_output = _randn( (token_count, hidden_size), seed=54, device=device, dtype=torch.bfloat16 @@ -420,9 +415,15 @@ def _tp1_vs_tpn_train_infer_worker(rank, world_size, init_method, result_queue, assert train_match, f"TP=1 vs TP={world_size} train forward mismatch" assert hidden_match, f"TP=1 vs TP={world_size} hidden grad mismatch" else: - assert not infer_match, f"TP=1 vs TP={world_size} infer forward unexpectedly matched" - assert not train_match, f"TP=1 vs TP={world_size} train forward unexpectedly matched" - assert not hidden_match, f"TP=1 vs TP={world_size} hidden grad unexpectedly matched" + assert not infer_match, ( + f"TP=1 vs TP={world_size} infer forward unexpectedly matched" + ) + assert not train_match, ( + f"TP=1 vs TP={world_size} train forward unexpectedly matched" + ) + assert not hidden_match, ( + f"TP=1 vs TP={world_size} hidden grad unexpectedly matched" + ) assert torch.equal( tp1_inputs[1].grad[feat_start:feat_end], tpn_inputs[1].grad @@ -456,7 +457,9 @@ def _make_topology_inputs(token_count, device): def _canonical(hidden, gate, up, down, grad): with torch.no_grad(): infer = qwen3_ffn(hidden, gate, up, down) - inputs = [value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down)] + inputs = [ + value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down) + ] train = qwen3_ffn(*inputs) train.backward(grad) return infer, train, inputs @@ -579,7 +582,9 @@ def _topology_worker(rank, world_size, init_method, result_queue, configs): if token_count not in canonical: tensors = _make_topology_inputs(token_count, device) canonical[token_count] = (*tensors, *_canonical(*tensors)) - hidden, gate, up, down, grad, infer_ref, train_ref, ref_inputs = canonical[token_count] + hidden, gate, up, down, grad, infer_ref, train_ref, ref_inputs = canonical[ + token_count + ] _run_topology_config( rank, dist, @@ -633,7 +638,7 @@ def _cache_worker(rank, world_size, init_method, result_queue): small = _ffn_tensors(8, device, seed=100, intermediate=128) first = qwen3_ffn(*small, tp_group=dist.group.WORLD) assert len(ffn_module._COLLECTIVES) == 1 - ((cache_key, first_collective),) = ffn_module._COLLECTIVES.items() + (cache_key, first_collective), = ffn_module._COLLECTIVES.items() first_handle = first_collective._handle first_capacity = first_collective.max_size_bytes assert first_handle != 0 @@ -706,11 +711,7 @@ def _uneven_sp_worker(rank, world_size, init_method, result_queue): sequence_parallel=True, ) result_queue.put( - { - "ok": False, - "rank": rank, - "failures": "uneven SP tokens should have raised", - } + {"ok": False, "rank": rank, "failures": "uneven SP tokens should have raised"} ) except ValueError as exc: message = str(exc) @@ -727,7 +728,9 @@ def _uneven_sp_worker(rank, world_size, init_method, result_queue): def _qwen3_8b_weights(device): - hidden = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=90, device=device, dtype=torch.bfloat16) + hidden = _randn( + (8, QWEN3_8B_HIDDEN_SIZE), seed=90, device=device, dtype=torch.bfloat16 + ) gate = _randn( (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), seed=91, @@ -746,7 +749,9 @@ def _qwen3_8b_weights(device): device=device, dtype=torch.bfloat16, ) - grad = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=94, device=device, dtype=torch.bfloat16) + grad = _randn( + (8, QWEN3_8B_HIDDEN_SIZE), seed=94, device=device, dtype=torch.bfloat16 + ) return hidden, gate, up, down, grad @@ -766,7 +771,8 @@ def _qwen3_8b_tp2_worker(rank, world_size, init_method, result_queue): with torch.no_grad(): infer_tp1 = qwen3_ffn(hidden, gate, up, down) tp1_inputs = [ - value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down) + value.detach().clone().requires_grad_(True) + for value in (hidden, gate, up, down) ] train_tp1 = qwen3_ffn(*tp1_inputs) train_tp1.backward(grad) @@ -843,86 +849,6 @@ def test_qwen_ffn_backward_matches_autograd_reference(monkeypatch): assert stub.calls.count("swiglu_backward") == 1 -def test_qwen_ffn_disable_split_k_false_uses_torch_matmul(monkeypatch): - stub = _TorchKernelStub() - monkeypatch.setattr(ffn_module, "_C", stub) - monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) - monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) - - hidden = _randn((2, 3, 8), seed=0) - gate_weight = _randn((12, 8), seed=1) - up_weight = _randn((12, 8), seed=2) - down_weight = _randn((8, 12), seed=3) - grad_output = _randn(hidden.shape, seed=4) - - ref_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - expected, _, _, _ = _reference(*ref_inputs) - expected.backward(grad_output) - - actual_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - actual = qwen3_ffn(*actual_inputs, disable_split_k=False) - actual.backward(grad_output) - - torch.testing.assert_close(actual, expected.detach()) - for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): - torch.testing.assert_close(actual_input.grad, reference.grad) - - assert stub.calls.count("det_gemm_fwd") == 0 - assert stub.calls.count("det_gemm_db") == 0 - assert stub.calls.count("swiglu_forward") == 1 - assert stub.calls.count("swiglu_backward") == 1 - - -def test_qwen_ffn_deterministic_false_uses_production_gemm(monkeypatch): - modes = [] - monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) - monkeypatch.setattr(ffn_module, "_require_ffn_kernels", lambda **kwargs: modes.append(kwargs)) - monkeypatch.setattr( - ffn_module._DeterministicFFNFunction, - "apply", - lambda *args: args[-1], - ) - - tensors = [torch.empty(1)] * 4 - assert qwen3_ffn(*tensors, deterministic=False) is False - assert modes == [{"disable_split_k": False}] - - -def test_qwen_ffn_rejects_conflicting_backend_switches(): - tensors = [torch.empty(1)] * 4 - - with pytest.raises(ValueError, match="conflicting FFN backends"): - qwen3_ffn(*tensors, deterministic=True, disable_split_k=False) - - -def test_qwen_ffn_rejects_non_bool_deterministic(): - tensors = [torch.empty(1)] * 4 - - with pytest.raises(TypeError, match="deterministic must be a bool or None"): - qwen3_ffn(*tensors, deterministic=1) # type: ignore[arg-type] - - -def test_qwen_ffn_rejects_non_bool_disable_split_k(): - hidden = torch.empty((2, 8), dtype=torch.bfloat16) - gate_weight = torch.empty((12, 8), dtype=torch.bfloat16) - up_weight = torch.empty((12, 8), dtype=torch.bfloat16) - down_weight = torch.empty((8, 12), dtype=torch.bfloat16) - with pytest.raises(TypeError, match="disable_split_k must be a bool"): - qwen3_ffn( - hidden, - gate_weight, - up_weight, - down_weight, - disable_split_k=1, # type: ignore[arg-type] - ) - - def test_qwen_ffn_rejects_non_huggingface_weight_layout(): hidden = torch.empty((2, 8), dtype=torch.bfloat16) gate_weight = torch.empty((8, 12), dtype=torch.bfloat16) @@ -934,8 +860,7 @@ def test_qwen_ffn_rejects_non_huggingface_weight_layout(): @requires_cuda_ffn -@pytest.mark.parametrize("disable_split_k", [True, False]) -def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(disable_split_k): +def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(): hidden = _randn((2, 3, 64), seed=10, device="cuda", dtype=torch.bfloat16) gate_weight = _randn((128, 64), seed=11, device="cuda", dtype=torch.bfloat16) up_weight = _randn((128, 64), seed=12, device="cuda", dtype=torch.bfloat16) @@ -953,7 +878,7 @@ def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(disable_split_k): value.detach().clone().requires_grad_(True) for value in (hidden, gate_weight, up_weight, down_weight) ] - actual = qwen3_ffn(*actual_inputs, disable_split_k=disable_split_k) + actual = qwen3_ffn(*actual_inputs) actual.backward(grad_output) torch.testing.assert_close(