Organize FE OSS tests by feature

docs/fe-oss-apis/attention/sdpa_bwd_d256.md

@@ -275,6 +275,6 @@ Tuple unpacking order is: `(dq_tensor, dk_tensor, dv_tensor)`.
 
 For runnable examples and reference-comparison checks, see:
 
-- `test/python/fe_api/test_sdpa_bwd.py`
-- `test/python/fe_api/test_sdpa_bwd_utils.py`
+- `test/python/fe_api/sdpa/test_sdpa_bwd.py`
+- `test/python/fe_api/sdpa/test_sdpa_bwd_utils.py`
 

docs/fe-oss-apis/bsa.md

@@ -214,7 +214,7 @@ The current public surface consists of allocating function wrappers under
 lifecycle for BSA.
 
 Correctness tests and FP32 references are under
-`test/python/fe_api/block_sparse_attention`.
+`test/python/fe_api/bsa`.
 
 ## Source provenance
 

docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md

@@ -465,5 +465,5 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 ## Usage Examples
 
 For runnable examples and validation, see:
-- `test/python/fe_api/test_discrete_grouped_gemm_dswiglu.py`
-- `test/python/fe_api/test_discrete_grouped_gemm_dswiglu_utils.py`
+- `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py`
+- `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py`

docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md

@@ -424,5 +424,5 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 ## Usage Examples
 
 For runnable examples and validation, see:
-- `test/python/fe_api/test_discrete_grouped_gemm_swiglu.py`
-- `test/python/fe_api/test_discrete_grouped_gemm_swiglu_utils.py`
+- `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu.py`
+- `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py`

docs/fe-oss-apis/gemm_fusions/gemm_amax.md

@@ -204,4 +204,4 @@ Tuple unpacking order is: `(c_tensor, amax_tensor)`.
 
 ## Usage examples
 
-For usage examples, see test cases in `test/python/fe_api/test_gemm_amax.py`
+For usage examples, see test cases in `test/python/fe_api/gemm/test_gemm_amax.py`

docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md

@@ -249,5 +249,5 @@ Tuple unpacking order is: `(d_tensor, dprob_tensor, amax_tensor, sfd_tensor)`.
 
 For end-to-end usage and regression coverage, see:
 
-- `test/python/fe_api/test_gemm_dsrelu.py`
-- `test/python/fe_api/test_gemm_dsrelu_utils.py`
+- `test/python/fe_api/gemm/test_gemm_dsrelu.py`
+- `test/python/fe_api/gemm/test_gemm_dsrelu_utils.py`

docs/fe-oss-apis/gemm_fusions/gemm_srelu.md

@@ -234,5 +234,5 @@ Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, sfd_tensor)`.
 
 For end-to-end usage and regression coverage, see:
 
-- `test/python/fe_api/test_gemm_srelu.py`
-- `test/python/fe_api/test_gemm_srelu_utils.py`
+- `test/python/fe_api/gemm/test_gemm_srelu.py`
+- `test/python/fe_api/gemm/test_gemm_srelu_utils.py`

docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md

@@ -350,4 +350,4 @@ Additional constraints:
 
 ## Usage examples
 
-For usage examples, see test cases in `test/python/fe_api/test_gemm_swiglu.py`
+For usage examples, see test cases in `test/python/fe_api/gemm/test_gemm_swiglu.py`

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md

@@ -396,4 +396,4 @@ Returns a `TupleDict` (dictionary + tuple unpacking):
 
 ## Usage Examples
 
-For usage examples, see test cases in `test/python/fe_api/test_grouped_gemm_dglu.py` (dense mode, unified API) and `test/python/fe_api/test_discrete_grouped_gemm_dswiglu.py` (discrete mode).
+For usage examples, see test cases in `test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py` (dense mode, unified API) and `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py` (discrete mode).

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md

@@ -314,5 +314,5 @@ Tuple unpacking order is: `(d_row_tensor, d_col_tensor, dprob_tensor, dbias_tens
 
 For end-to-end usage and regression coverage, see:
 
-- `test/python/fe_api/test_grouped_gemm_dsrelu.py`
-- `test/python/fe_api/test_grouped_gemm_dsrelu_utils.py`
+- `test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py`
+- `test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_utils.py`

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md

@@ -410,4 +410,4 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 
 ## Usage Examples
 
-For usage examples, see test cases in `test/python/fe_api/test_grouped_gemm_dswiglu.py` + `test/python/fe_api/test_grouped_gemm_dswiglu_utils.py`
+For usage examples, see test cases in `test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu.py` + `test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py`

docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md

@@ -389,4 +389,4 @@ Returns a `TupleDict` (dictionary + tuple unpacking):
 
 ## Usage Examples
 
-For usage examples, see test cases in `test/python/fe_api/test_grouped_gemm_glu.py` (dense mode, unified API) and `test/python/fe_api/test_discrete_grouped_gemm_swiglu.py` (discrete mode).
+For usage examples, see test cases in `test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py` (dense mode, unified API) and `test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu.py` (discrete mode).

docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md

@@ -372,4 +372,4 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 
 ## Usage Examples
 
-For usage examples, see test cases in `test/python/fe_api/test_grouped_gemm_quant.py` + `test/python/fe_api/test_grouped_gemm_quant_utils.py`
+For usage examples, see test cases in `test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py` + `test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py`

docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant_unified.md

@@ -244,4 +244,4 @@ Returns `TupleDict`: `d_tensor`, `d_col_tensor` (optional; `None` for `bfloat16`
 
 ## Usage Examples
 
-For usage examples, see `test/python/fe_api/test_grouped_gemm_quant.py` + `test/python/fe_api/test_grouped_gemm_quant_utils.py` (dense and discrete unified API coverage)
+For usage examples, see `test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py` + `test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py` (dense and discrete unified API coverage)

docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md

@@ -298,5 +298,5 @@ Tuple unpacking order is: `(c_tensor, d_tensor, d_col_tensor, amax_tensor, sfd_r
 
 For end-to-end usage and regression coverage, see:
 
-- `test/python/fe_api/test_grouped_gemm_srelu.py`
-- `test/python/fe_api/test_grouped_gemm_srelu_utils.py`
+- `test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu.py`
+- `test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_utils.py`

docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md

@@ -386,4 +386,4 @@ Returns a `TupleDict` - a dictionary-like object that also supports tuple unpack
 
 ## Usage Examples
 
-For usage examples, see test cases in `test/python/fe_api/test_grouped_gemm_swiglu.py` + `test/python/fe_api/test_grouped_gemm_swiglu_utils.py`
+For usage examples, see test cases in `test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu.py` + `test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_utils.py`

docs/fe-oss-apis/rmsnorm_rht_amax.md

@@ -148,4 +148,4 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`.
 ## Verification
 
 Focused correctness and cache coverage live in:
-- `test/python/fe_api/test_rmsnorm_rht_amax.py`
+- `test/python/fe_api/norm/test_rmsnorm_rht_amax.py`

python/cudnn/native_sparse_attention/sparse_attention.md

@@ -36,10 +36,10 @@ pip install nvidia-cudnn-frontend[cutedsl]
 ## Usage
 
 Sample usage and tests can be found in the (test/python) folder:
-- [test_NSA_selection_attention.py](test/python/fe_api/nsa/test_NSA_selection_attention.py), `pytest test/python/test_NSA_selection_attention.py`
-- [test_NSA_topk_reduction.py](test/python/fe_api/nsa/test_NSA_topk_reduction.py), `pytest test/python/test_NSA_topk_reduction.py`
-- [test_NSA_compression_attention.py](test/python/fe_api/nsa/test_NSA_compression_attention.py), `pytest test/python/test_NSA_compression_attention.py`
-- [test_NSA_swa.py](test/python/fe_api/nsa/test_NSA_swa.py), `pytest test/python/test_NSA_swa.py`
+- [test_NSA_selection_attention.py](test/python/fe_api/nsa/test_NSA_selection_attention.py), `pytest test/python/fe_api/nsa/test_NSA_selection_attention.py`
+- [test_NSA_topk_reduction.py](test/python/fe_api/nsa/test_NSA_topk_reduction.py), `pytest test/python/fe_api/nsa/test_NSA_topk_reduction.py`
+- [test_NSA_compression_attention.py](test/python/fe_api/nsa/test_NSA_compression_attention.py), `pytest test/python/fe_api/nsa/test_NSA_compression_attention.py`
+- [test_NSA_swa.py](test/python/fe_api/nsa/test_NSA_swa.py), `pytest test/python/fe_api/nsa/test_NSA_swa.py`
 
 Once all components are implemented, we will offer a central NSA API that will do the full NSA computation end-to-end. We will also offer the individual components as standalone APIs, as demonstrated below.
 

skills/cutedsl-kernel-integration/references/integration-pattern.md

@@ -136,9 +136,9 @@ Add focused pytest coverage under `test/python/fe_api/`.
 
 Typical files:
 
-- `test/python/fe_api/test_<operation>.py`
-- `test/python/fe_api/test_<operation>_utils.py` for reusable test helpers or shape/reference utilities.
-- A family subdirectory when matching existing structure, such as `test/python/fe_api/nsa/`.
+- `test/python/fe_api/<feature>/test_<operation>.py`
+- `test/python/fe_api/<feature>/test_<operation>_utils.py` for reusable test helpers or shape/reference utilities.
+- A nested family subdirectory when matching existing structure, such as `test/python/fe_api/nsa/`.
 
 Coverage should include:
 

test/python/fe_api/README.md

@@ -0,0 +1,11 @@
+# FE OSS tests
+
+Tests are organized by feature:
+
+- `gemm/`: dense GEMM fusions
+- `grouped_gemm/`: grouped and discrete GEMM fusions
+- `sdpa/`: scaled dot-product attention
+- `bsa/`: Block Sparse Attention
+- `dsa/`: DeepSeek Sparse Attention
+- `nsa/`: Native Sparse Attention
+- `norm/`: normalization fusions

test/python/fe_api/block_sparse_attention/test_BSA_attention_backward.py → test/python/fe_api/bsa/test_BSA_attention_backward.py

@@ -7,8 +7,8 @@
 import torch
 
 from test_utils import torch_fork_set_rng
-from fe_api.block_sparse_attention.bsa_reference import attention_backward_reference, block_sparse_mask
-from fe_api.block_sparse_attention.bsa_utils import make_fixed_metadata, supported_block_size
+from fe_api.bsa.bsa_reference import attention_backward_reference, block_sparse_mask
+from fe_api.bsa.bsa_utils import make_fixed_metadata, supported_block_size
 
 pytestmark = [pytest.mark.gpu_exclusive, pytest.mark.xdist_group(name="gpu_exclusive")]
 

test/python/fe_api/block_sparse_attention/test_BSA_attention_forward.py → test/python/fe_api/bsa/test_BSA_attention_forward.py

@@ -8,8 +8,8 @@
 import torch
 
 from test_utils import torch_fork_set_rng
-from fe_api.block_sparse_attention.bsa_reference import attention_reference, block_sparse_mask
-from fe_api.block_sparse_attention.bsa_utils import make_fixed_metadata, make_variable_metadata, supported_block_size
+from fe_api.bsa.bsa_reference import attention_reference, block_sparse_mask
+from fe_api.bsa.bsa_utils import make_fixed_metadata, make_variable_metadata, supported_block_size
 
 pytestmark = [pytest.mark.gpu_exclusive, pytest.mark.xdist_group(name="gpu_exclusive")]
 

test/python/fe_api/test_gemm_amax.py → test/python/fe_api/gemm/test_gemm_amax.py

@@ -3,7 +3,7 @@
 import pytest
 
 from test_utils import torch_fork_set_rng
-from fe_api.test_gemm_amax_utils import (
+from fe_api.gemm.test_gemm_amax_utils import (
     with_gemm_amax_params_fp4,
     with_gemm_amax_params_fp8,
 )
@@ -155,7 +155,7 @@ def _test_gemm_amax_compile_execute(
     try:
         from cudnn import GemmAmaxSm100
         from cuda.bindings import driver as cuda
-        from fe_api.test_gemm_amax_utils import (
+        from fe_api.gemm.test_gemm_amax_utils import (
             allocate_input_tensors,
             allocate_output_tensors,
             check_ref_gemm_amax,
@@ -242,7 +242,7 @@ def _test_gemm_amax_wrapper(
     try:
         from cudnn import gemm_amax_wrapper_sm100
         from cuda.bindings import driver as cuda
-        from fe_api.test_gemm_amax_utils import (
+        from fe_api.gemm.test_gemm_amax_utils import (
             allocate_input_tensors,
             allocate_output_tensors,
             check_ref_gemm_amax,

test/python/fe_api/test_gemm_amax_utils.py → test/python/fe_api/gemm/test_gemm_amax_utils.py

@@ -10,7 +10,7 @@
     _bfloat16_to_float4_e2m1fn_x2,
     float4_e2m1fn_x2_to_float32,
 )
-from test_fe_api_utils import create_and_permute_tensor, create_scale_factor_tensor
+from fe_api.test_fe_api_utils import create_and_permute_tensor, create_scale_factor_tensor
 
 # Parameterization marks for GEMM Amax
 GEMM_AMAX_PARAM_MARKS_FP4 = [

test/python/fe_api/test_gemm_dsrelu.py → test/python/fe_api/gemm/test_gemm_dsrelu.py

@@ -4,7 +4,7 @@
 import cudnn
 from test_utils import torch_fork_set_rng
 
-from fe_api.test_gemm_dsrelu_utils import (
+from fe_api.gemm.test_gemm_dsrelu_utils import (
     allocate_gemm_dsrelu_outputs,
     allocate_gemm_dsrelu_tensors,
     check_ref_gemm_dsrelu,

test/python/fe_api/test_gemm_dsrelu_utils.py → test/python/fe_api/gemm/test_gemm_dsrelu_utils.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,

test/python/fe_api/test_gemm_srelu.py → test/python/fe_api/gemm/test_gemm_srelu.py

@@ -4,7 +4,7 @@
 import cudnn
 from test_utils import torch_fork_set_rng
 
-from fe_api.test_gemm_srelu_utils import (
+from fe_api.gemm.test_gemm_srelu_utils import (
     allocate_gemm_srelu_outputs,
     allocate_gemm_srelu_tensors,
     check_ref_gemm_srelu,

test/python/fe_api/test_gemm_srelu_utils.py → test/python/fe_api/gemm/test_gemm_srelu_utils.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,

test/python/fe_api/test_gemm_swiglu.py → test/python/fe_api/gemm/test_gemm_swiglu.py

@@ -2,7 +2,7 @@
 
 import pytest
 from test_utils import torch_fork_set_rng
-from fe_api.test_gemm_swiglu_utils import (
+from fe_api.gemm.test_gemm_swiglu_utils import (
     allocate_input_tensors,
     allocate_output_tensors,
     check_ref_gemm_swiglu,

test/python/fe_api/test_gemm_swiglu_utils.py → test/python/fe_api/gemm/test_gemm_swiglu_utils.py

@@ -6,7 +6,7 @@
 import torch
 import pytest
 from typing import Optional, Tuple
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,

test/python/fe_api/test_discrete_grouped_gemm_dswiglu.py → test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py

@@ -6,7 +6,7 @@
 import pytest
 from test_utils import torch_fork_set_rng
 from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES
-from fe_api.test_discrete_grouped_gemm_dswiglu_utils import (
+from fe_api.grouped_gemm.test_discrete_grouped_gemm_dswiglu_utils import (
     discrete_dswiglu_init,
     with_discrete_dswiglu_params_fp4,
     with_discrete_dswiglu_params_fp8,

test/python/fe_api/test_discrete_grouped_gemm_dswiglu_utils.py → test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_utils.py

@@ -5,12 +5,12 @@
 import torch
 import pytest
 from typing import Tuple, List, Dict, Any, Optional
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,
 )
-from fe_api.test_grouped_gemm_dswiglu_utils import compute_reference_row_quant
+from fe_api.grouped_gemm.test_grouped_gemm_dswiglu_utils import compute_reference_row_quant
 
 # =============================================================================
 # Parameterization Marks

test/python/fe_api/test_discrete_grouped_gemm_swiglu.py → test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu.py

@@ -9,7 +9,7 @@
 import pytest
 from test_utils import torch_fork_set_rng
 from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES
-from fe_api.test_discrete_grouped_gemm_swiglu_utils import (
+from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import (
     discrete_grouped_gemm_init,
     with_discrete_grouped_gemm_params_fp4,
     with_discrete_grouped_gemm_params_fp8,
@@ -466,8 +466,8 @@ def test_discrete_vs_contiguous_match(ab_dtype, d_dtype, sf_vec_size, sf_dtype,
     if major < 10:
         pytest.skip(f"Requires SM100+, found SM{major}0")
 
-    from test_fe_api_utils import create_and_permute_tensor, create_scale_factor_tensor
-    from fe_api.test_discrete_grouped_gemm_swiglu_utils import create_mask
+    from fe_api.test_fe_api_utils import create_and_permute_tensor, create_scale_factor_tensor
+    from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import create_mask
     from cudnn.api_base import ceil_div
 
     n, k, num_experts = 512, 512, 4

test/python/fe_api/test_discrete_grouped_gemm_swiglu_utils.py → test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_utils.py

@@ -6,7 +6,7 @@
 import torch
 import pytest
 from typing import Tuple, List, Dict, Any, Optional
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,

test/python/fe_api/test_grouped_gemm_dglu.py → test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py

@@ -9,19 +9,19 @@
 import pytest
 from test_utils import torch_fork_set_rng
 from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES
-from fe_api.test_grouped_gemm_swiglu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import (
     grouped_gemm_swiglu_init,
     allocate_grouped_gemm_input_tensors as allocate_grouped_gemm_input_tensors_base,
 )
-from fe_api.test_grouped_gemm_dswiglu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_dswiglu_utils import (
     with_grouped_gemm_dswiglu_params_fp4,
     with_grouped_gemm_dswiglu_params_fp8,
     with_grouped_gemm_dswiglu_params_dbias_fp4,
     with_grouped_gemm_dswiglu_params_dbias_fp8,
     allocate_grouped_gemm_dswiglu_tensors,
     check_ref_grouped_gemm_dswiglu,
 )
-from fe_api.test_discrete_grouped_gemm_dswiglu_utils import (
+from fe_api.grouped_gemm.test_discrete_grouped_gemm_dswiglu_utils import (
     discrete_dswiglu_init,
     allocate_discrete_dswiglu_input_tensors,
     allocate_discrete_dswiglu_output_tensors,

test/python/fe_api/test_grouped_gemm_dsrelu.py → test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py

@@ -8,15 +8,15 @@
 import torch
 import pytest
 from test_utils import torch_fork_set_rng
-from fe_api.test_grouped_gemm_dsrelu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_dsrelu_utils import (
     with_grouped_gemm_dsrelu_params_fp4,
     with_grouped_gemm_dsrelu_params_fp8,
     allocate_grouped_gemm_dsrelu_tensors,
     allocate_grouped_gemm_input_tensors,
     check_ref_grouped_gemm_dsrelu,
     grouped_gemm_dsrelu_init,
 )
-from fe_api.test_discrete_grouped_gemm_swiglu_utils import (
+from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import (
     allocate_discrete_input_tensors,
     discrete_grouped_gemm_init,
 )

test/python/fe_api/test_grouped_gemm_dsrelu_utils.py → test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_utils.py

@@ -6,15 +6,15 @@
 import torch
 import pytest
 from typing import Optional, Tuple, List, Dict, Any
-from test_fe_api_utils import (
+from fe_api.test_fe_api_utils import (
     ceil_div,
     compute_reference_amax,
     create_and_permute_tensor,
     create_scale_factor_tensor,
     create_sf_layout_tensor,
     cvt_sf_MKL_to_M32x4xrm_K4xrk_L,
 )
-from test_grouped_gemm_swiglu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import (
     allocate_grouped_gemm_input_tensors as allocate_grouped_gemm_input_tensors_base,
     get_dtype_rcp_limits as get_grouped_gemm_dtype_rcp_limits,
     grouped_gemm_swiglu_init as grouped_gemm_dsrelu_init,

test/python/fe_api/test_grouped_gemm_dswiglu.py → test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu.py

@@ -8,11 +8,11 @@
 import torch
 import pytest
 from test_utils import torch_fork_set_rng
-from fe_api.test_grouped_gemm_swiglu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import (
     grouped_gemm_swiglu_init,
     allocate_grouped_gemm_input_tensors as allocate_grouped_gemm_input_tensors_base,
 )
-from fe_api.test_grouped_gemm_dswiglu_utils import (
+from fe_api.grouped_gemm.test_grouped_gemm_dswiglu_utils import (
     with_grouped_gemm_dswiglu_params_fp4,
     with_grouped_gemm_dswiglu_params_fp8,
     allocate_grouped_gemm_dswiglu_tensors,