# ruff: noqa: N803, E741 """Layer-level correctness & performance test for MoT parallel linear layers. Compares two equivalent computation paths: - Reference: 2x standard vLLM parallel linear layers + PyTorch index scatter/gather (text_linear(x[text_idx]) + vae_linear(x[vae_idx])) - Target: 1x MoT fused parallel linear layer (mot_linear(x, text_indices, vae_indices)) The reference path uses cuBLAS GEMM (always auto-tuned by cuBLAS). The MoT path uses a fused Triton kernel whose tile config is loaded from a JSON file matched by ``device + dtype``. If no tuned config is found for the current GPU, the kernel falls back to a conservative default and a warning is printed — the correctness test still passes but the performance comparison is NOT representative. Usage:: pytest tests/diffusion/kernels/mot/test_mot_linear.py -v -s """ from __future__ import annotations import os import time from dataclasses import dataclass import pytest import torch from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.linear import ( QKVParallelLinear, RowParallelLinear, ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, init_distributed_environment, initialize_model_parallel, model_parallel_is_initialized, ) from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import ( MoTQKVParallelLinear, ) from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import ( MoTRowParallelLinear, ) from vllm_omni.diffusion.layers.mot.ops.mot_gemm import ( get_best_mot_config, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- # BAGEL-7B-MoT architecture parameters _BAGEL_HEAD_SIZE = 128 _BAGEL_TOTAL_NUM_HEADS = 28 _BAGEL_TOTAL_NUM_KV_HEADS = 4 _VAE_CHUNK_SIZE = 1024 # the token number of one image _IMAGE_NUM = [1, 2, 4, 8] @pytest.fixture(scope="module", autouse=True) def _init_single_rank_tp_env(): """Initialize single-rank distributed/TP env for vLLM linear params.""" os.environ.setdefault("RANK", "0") os.environ.setdefault("LOCAL_RANK", "0") os.environ.setdefault("WORLD_SIZE", "1") os.environ.setdefault("MASTER_ADDR", "127.0.0.1") os.environ.setdefault("MASTER_PORT", "29501") if not torch.distributed.is_initialized(): init_distributed_environment(world_size=1, rank=0, local_rank=0) if not model_parallel_is_initialized(): initialize_model_parallel( data_parallel_size=1, cfg_parallel_size=1, sequence_parallel_size=1, ulysses_degree=1, ring_degree=1, tensor_parallel_size=1, pipeline_parallel_size=1, ) yield destroy_distributed_env() # --------------------------------------------------------------------------- # DType configuration — extensible for future quantized types # --------------------------------------------------------------------------- @dataclass class DTypeConfig: """Parsed dtype configuration for a test case.""" torch_dtype: torch.dtype use_fp8_w8a8: bool = False use_int8_w8a16: bool = False use_int4_w4a16: bool = False def _parse_dtype(dtype_str: str) -> DTypeConfig: """Parse a dtype string into quantization flags and torch dtype. Supported now: "w16a16_bf16" — BF16 weights & activations (no quantization) "w16a16_fp16" — FP16 weights & activations (no quantization) Reserved for future: "fp8_w8a8" — FP8 W8A8 quantization "int8_w8a16" — INT8 weight-only quantization "int4_w4a16" — INT4 weight-only quantization """ supported: dict[str, DTypeConfig] = { "w16a16_bf16": DTypeConfig(torch_dtype=torch.bfloat16), "w16a16_fp16": DTypeConfig(torch_dtype=torch.float16), } if dtype_str in supported: return supported[dtype_str] pytest.skip(f"Quantized dtype '{dtype_str}' not yet implemented in layer test") # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _report_mot_config(K: int, N: int, M: int): """Print which Triton tile config the MoT kernel will use.""" loaded_m_key, config = get_best_mot_config(M, N, K, None) if loaded_m_key == -1: print( " [config] WARNING: No tuned config found — " "using conservative default. " "Performance numbers are NOT representative. " "Run mot_linear_benchmarks.py --tune to generate configs." ) else: print(f" [config] Tuned config loaded (actual M={M}, loaded M={loaded_m_key}) config = {config})") def _make_indices(image_num: int, vae_chunk_size: int, device: str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]: """ Simulate exact Bagel-MoT distributions for image generation: Pattern per image like: [1 Text] + [4096 VAE] + [1 Text] Returns text_indices, vae_indices, and the exact total M. """ text_idx_list = [] vae_idx_list = [] current_idx = 0 for _ in range(image_num): text_idx_list.append(current_idx) current_idx += 1 vae_idx_list.extend(range(current_idx, current_idx + vae_chunk_size)) current_idx += vae_chunk_size text_idx_list.append(current_idx) current_idx += 1 text_indices = torch.tensor(text_idx_list, dtype=torch.long, device=device) vae_indices = torch.tensor(vae_idx_list, dtype=torch.long, device=device) exact_M = current_idx # exact_M = image_num * (vae_chunk_size + 2) return text_indices, vae_indices, exact_M def _benchmark(fn, warmup: int = 20, iters: int = 100) -> float: """Return mean latency in milliseconds.""" cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda") for _ in range(warmup): fn() torch.accelerator.synchronize() t_flush_start = time.perf_counter() for _ in range(iters): cache_flusher.zero_() torch.accelerator.synchronize() flush_time_total = time.perf_counter() - t_flush_start # 3. Measure the total time of "flush cache + operator execution" torch.accelerator.synchronize() t_total_start = time.perf_counter() for _ in range(iters): cache_flusher.zero_() fn() torch.accelerator.synchronize() total_time = time.perf_counter() - t_total_start # 4. Asynchronous subtraction separation # The E2E time of the pure operator = total time - flush time pure_fn_time_total = total_time - flush_time_total avg_ms = (pure_fn_time_total / iters) * 1000.0 del cache_flusher return max(avg_ms, 0.001) def _sync_weights(ref_text, ref_vae, mot_layer): """Assign the same random weights to reference layers and MoT layer.""" with torch.no_grad(): W_text = torch.randn_like(ref_text.weight) * 0.02 W_vae = torch.randn_like(ref_vae.weight) * 0.02 ref_text.weight.copy_(W_text) ref_vae.weight.copy_(W_vae) mot_layer.weight.copy_(W_text) mot_layer.gen_exp.weight.copy_(W_vae) if ref_text.bias is not None and mot_layer.bias is not None: b_text = torch.randn_like(ref_text.bias) * 0.02 ref_text.bias.copy_(b_text) mot_layer.bias.copy_(b_text) if ref_vae.bias is not None and mot_layer.gen_exp.bias is not None: b_vae = torch.randn_like(ref_vae.bias) * 0.02 ref_vae.bias.copy_(b_vae) mot_layer.gen_exp.bias.copy_(b_vae) def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear): """Reference path: index-gather → 2x standard linear → index-scatter.""" M = x.size(0) out_text = text_linear(x[text_indices]) out_vae = vae_linear(x[vae_indices]) if isinstance(out_text, tuple): out_text = out_text[0] if isinstance(out_vae, tuple): out_vae = out_vae[0] N = out_text.size(-1) output = torch.empty(M, N, dtype=x.dtype, device=x.device) output[text_indices] = out_text output[vae_indices] = out_vae return output def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str): """Compare outputs, print metrics, and assert correctness. Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16). We upcast to fp32 solely for computing error metrics with higher arithmetic precision — the actual layer outputs remain bf16. """ # Upcast for numerically stable error computation only ref_hp = ref.float() mot_hp = mot.float() abs_err = (ref_hp - mot_hp).abs() max_abs = abs_err.max().item() # Mixed metric: relative error where |ref| >= 1, absolute error otherwise denom = ref_hp.abs().clamp(min=1.0) max_rel = (abs_err / denom).max().item() cos_sim = ( torch.nn.functional.cosine_similarity( ref_hp, mot_hp, dim=-1, ) .min() .item() ) print(f"\n [{tag}] max_abs={max_abs:.4e} max_rel={max_rel:.4e} min_cos_sim={cos_sim:.6f}") # Cosine similarity is the primary correctness gate: robust to scale # and accumulation-order differences between cuBLAS and Triton. # For bf16 GEMM with K up to ~20k, cos_sim > 0.99 is easily achieved. assert cos_sim > 0.98, f"Cosine similarity too low: {cos_sim:.6f}" # Supplementary per-element check (generous to avoid flaky failures # on extreme K dimensions like 18944) assert max_rel < 0.1, f"Max relative error too large: {max_rel:.4e}" def _run_timing( ref_fn, mot_fn, tag: str, warmup: int = 20, iters: int = 100, ): """Benchmark both paths and print timing comparison.""" ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters) mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters) speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf") print(f" [{tag}] Ref(2x linear): {ref_ms:.3f} ms | MoT(fused): {mot_ms:.3f} ms | Speedup: {speedup:.2f}x") # ========================================================================= # Test: qkv proj # ========================================================================= @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 4608, "w16a16_fp16") for num in _IMAGE_NUM], ids=[f"img{num}_K3584_N4608_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N4608_fp16" for num in _IMAGE_NUM], ) def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool): dcfg = _parse_dtype(dtype) torch.manual_seed(42) with set_current_vllm_config(VllmConfig()): text_linear = QKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() vae_linear = QKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTQKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=bias, vae_bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() assert text_linear.output_size_per_partition == N, ( f"Expected output_size_per_partition={N}, " f"got {text_linear.output_size_per_partition}. " f"Check head parameters." ) _sync_weights(text_linear, vae_linear, mot_linear) text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") tag = f"QKVParallel M={M} K={K} N={N}" _report_mot_config(K, N, M) with torch.no_grad(): ref = _reference_forward( x, text_idx, vae_idx, text_linear, vae_linear, ) mot_out, _ = mot_linear(x, text_idx, vae_idx) _check_and_report(ref, mot_out, tag) with torch.no_grad(): _run_timing( lambda: _reference_forward( x, text_idx, vae_idx, text_linear, vae_linear, ), lambda: mot_linear(x, text_idx, vae_idx), tag, ) # ========================================================================= # Test: o proj # ========================================================================= @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 3584, "w16a16_fp16") for num in _IMAGE_NUM], ids=[f"img{num}_K3584_N3584_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N3584_fp16" for num in _IMAGE_NUM], ) def test_mot_o_proj( image_num: int, K: int, N: int, dtype: str, bias: bool, ): dcfg = _parse_dtype(dtype) torch.manual_seed(42) with set_current_vllm_config(VllmConfig()): text_linear = RowParallelLinear( K, N, bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() vae_linear = RowParallelLinear( K, N, bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTRowParallelLinear( K, N, bias=bias, vae_bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() _sync_weights(text_linear, vae_linear, mot_linear) text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") tag = f"O Proj M={M} K={K} N={N}" _report_mot_config(K, N, M) # Correctness (also warms up Triton JIT compilation) with torch.no_grad(): ref = _reference_forward( x, text_idx, vae_idx, text_linear, vae_linear, ) mot_out, _ = mot_linear(x, text_idx, vae_idx) _check_and_report(ref, mot_out, tag) # Performance with torch.no_grad(): _run_timing( lambda: _reference_forward( x, text_idx, vae_idx, text_linear, vae_linear, ), lambda: mot_linear(x, text_idx, vae_idx), tag, ) # ========================================================================= # Test: und-mode (text_indices=None) — falls back to standard forward # ========================================================================= @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "K, N, dtype", [ (3584, 4608, "w16a16_bf16"), (3584, 4608, "w16a16_fp16"), ], ids=["QKV_K3584_N4608_bf16", "QKV_K3584_N4608_fp16"], ) def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool): """und-mode: text_indices=None should produce same output as standard QKVParallelLinear.""" dcfg = _parse_dtype(dtype) torch.manual_seed(42) M = 1026 with set_current_vllm_config(VllmConfig()): ref_linear = QKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTQKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=bias, vae_bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() with torch.no_grad(): mot_linear.weight.copy_(ref_linear.weight) if bias and ref_linear.bias is not None: mot_linear.bias.copy_(ref_linear.bias) x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") with torch.no_grad(): ref_out, _ = ref_linear(x) mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None) _check_and_report(ref_out, mot_out, f"QKV und-mode M={M} K={K}") @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "K, N, dtype", [ (3584, 3584, "w16a16_bf16"), (3584, 3584, "w16a16_fp16"), ], ids=["Row_K3584_N3584_bf16", "Row_K3584_N3584_fp16"], ) def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool): """und-mode: text_indices=None should produce same output as standard RowParallelLinear.""" dcfg = _parse_dtype(dtype) torch.manual_seed(42) M = 1026 with set_current_vllm_config(VllmConfig()): ref_linear = RowParallelLinear( K, N, bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTRowParallelLinear( K, N, bias=bias, vae_bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() with torch.no_grad(): mot_linear.weight.copy_(ref_linear.weight) if bias and ref_linear.bias is not None: mot_linear.bias.copy_(ref_linear.bias) x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") with torch.no_grad(): ref_out, _ = ref_linear(x) mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None) _check_and_report(ref_out, mot_out, f"Row und-mode M={M} K={K}") # ========================================================================= # Test: boundary cases (all-text, all-VAE) # ========================================================================= @pytest.mark.parametrize( "boundary_mode", ["all_text", "all_vae"], ids=["all_text", "all_vae"], ) @pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"]) def test_mot_qkv_boundary(boundary_mode: str, dtype: str): """Boundary: all tokens routed to one expert.""" dcfg = _parse_dtype(dtype) torch.manual_seed(42) M = 512 K = 3584 with set_current_vllm_config(VllmConfig()): text_linear = QKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=False, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() vae_linear = QKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=False, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTQKVParallelLinear( hidden_size=K, head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, bias=False, vae_bias=False, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() _sync_weights(text_linear, vae_linear, mot_linear) all_indices = torch.arange(M, dtype=torch.long, device="cuda") empty_indices = torch.empty(0, dtype=torch.long, device="cuda") if boundary_mode == "all_text": text_idx, vae_idx = all_indices, empty_indices else: text_idx, vae_idx = empty_indices, all_indices x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") with torch.no_grad(): ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear) mot_out, _ = mot_linear(x, text_idx, vae_idx) _check_and_report(ref, mot_out, f"QKV boundary={boundary_mode} M={M}") @pytest.mark.parametrize( "boundary_mode", ["all_text", "all_vae"], ids=["all_text", "all_vae"], ) @pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"]) def test_mot_row_boundary(boundary_mode: str, dtype: str): """Boundary: all tokens routed to one expert for RowParallel.""" dcfg = _parse_dtype(dtype) torch.manual_seed(42) M = 512 K = 3584 N = 3584 with set_current_vllm_config(VllmConfig()): text_linear = RowParallelLinear( K, N, bias=False, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() vae_linear = RowParallelLinear( K, N, bias=False, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() mot_linear = MoTRowParallelLinear( K, N, bias=False, vae_bias=False, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() _sync_weights(text_linear, vae_linear, mot_linear) all_indices = torch.arange(M, dtype=torch.long, device="cuda") empty_indices = torch.empty(0, dtype=torch.long, device="cuda") if boundary_mode == "all_text": text_idx, vae_idx = all_indices, empty_indices else: text_idx, vae_idx = empty_indices, all_indices x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") with torch.no_grad(): ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear) mot_out, _ = mot_linear(x, text_idx, vae_idx) _check_and_report(ref, mot_out, f"Row boundary={boundary_mode} M={M}")