diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index 0ae55dfa7..e48dde4d4 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/tests/diffusion/kernels/mot/test_mot_linear.py @@ -295,12 +295,13 @@ def _run_timing( # ========================================================================= +@pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM], ids=lambda val: "", ) -def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str): +def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool): dcfg = _parse_dtype(dtype) torch.manual_seed(42) @@ -310,7 +311,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str): head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, - bias=True, + bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() @@ -319,7 +320,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str): head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, - bias=True, + bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() @@ -328,8 +329,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str): head_size=_BAGEL_HEAD_SIZE, total_num_heads=_BAGEL_TOTAL_NUM_HEADS, total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, - bias=True, - vae_bias=True, + bias=bias, + vae_bias=bias, params_dtype=dcfg.torch_dtype, disable_tp=True, ).cuda() @@ -379,6 +380,7 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str): # ========================================================================= +@pytest.mark.parametrize("bias", [True, False]) @pytest.mark.parametrize( "image_num, K, N, dtype", [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM], @@ -389,6 +391,7 @@ def test_mot_o_proj( K: int, N: int, dtype: str, + bias: bool, ): dcfg = _parse_dtype(dtype) torch.manual_seed(42) @@ -397,7 +400,7 @@ def test_mot_o_proj( text_linear = RowParallelLinear( K, N, - bias=False, + bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, @@ -405,7 +408,7 @@ def test_mot_o_proj( vae_linear = RowParallelLinear( K, N, - bias=False, + bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, @@ -413,7 +416,8 @@ def test_mot_o_proj( mot_linear = MoTRowParallelLinear( K, N, - bias=False, + bias=bias, + vae_bias=bias, input_is_parallel=True, params_dtype=dcfg.torch_dtype, disable_tp=True, diff --git a/vllm_omni/diffusion/layers/mot/configs/README b/vllm_omni/diffusion/layers/mot/configs/README new file mode 100644 index 000000000..6661c6c9a --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/configs/README @@ -0,0 +1,50 @@ +This directory contains auto-tuned Triton kernel configurations for the +MoT (Mixture-of-Tokens) GEMM and RMSNorm operators used by BAGEL and other +MoT-architecture diffusion models. + +File naming convention: + device_name=,dtype=.json + +For example: + device_name=NVIDIA_A100-SXM4-80GB,dtype=w16a16.json + +Each JSON file maps (K, N) matrix shapes to a dictionary of batch sizes (M) +and their optimal Triton tile configurations: + + { + "3584_4608": { // K=3584, N=4608 (QKV projection) + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + ... + }, + "3584_3584": { ... } // K=3584, N=3584 (output projection) + } + +Config loading order (3-tier, see ops/mot_gemm.py): + 1. $VLLM_TUNED_CONFIG_FOLDER/ (env override) + 2. This directory: vllm_omni/diffusion/layers/mot/configs/ + 3. Conservative default config (compiles everywhere, sub-optimal perf) + +If no config file matches the current device, a warning is printed with +instructions to run the auto-tuning benchmark. + +To generate configs for your hardware: + + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 1 --dtype w16a16 --tune \ + --save-dir vllm_omni/diffusion/layers/mot/configs/ + +For multi-GPU tuning (uses Ray for parallel search): + + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 2 --tune + +See benchmarks/kernels/mot_linear_benchmarks.py for full options. diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index f71526016..8915ecca5 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -396,60 +396,45 @@ class PackedAttentionMoT(nn.Module): ) -> tuple[torch.Tensor, NaiveCache | None]: """SP-aware attention for gen mode denoising. - Converts packed format to batched (1, S, H, D) and uses the diffusion - Attention layer (Ulysses / Ring) with joint mechanism: + Uses MoT unified layers for projection/norm (fused text/vae routing), + then splits into text/vae for the DiffusionAttention (Ulysses / Ring) + with joint mechanism: - Main Q/K/V: VAE tokens (split across SP ranks) - Joint Q: text marker Q (replicated) - Joint K/V: KV cache K/V + text marker K/V (replicated) """ + text_indices = packed_text_indexes + vae_indices = packed_vae_token_indexes + packed_query_sequence = packed_query_sequence.to(torch.bfloat16) - packed_text_query_sequence = packed_query_sequence[packed_text_indexes] - packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes] + # MoT QKV projection — fused kernel routes text/vae to correct weights + qkv, _ = self.qkv_proj(packed_query_sequence, text_indices, vae_indices) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q = q.view(-1, self.num_heads, self.head_dim) + k = k.view(-1, self.num_kv_heads, self.head_dim) + v = v.view(-1, self.num_kv_heads, self.head_dim) - # Project text tokens through base qkv - text_qkv, _ = self.qkv_proj(packed_text_query_sequence) - text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # MoT QK norms — routes text/vae to weight/gen_weight internally + q = self.q_norm(q.to(torch.float32), text_indices, vae_indices) + k = self.k_norm(k.to(torch.float32), text_indices, vae_indices) - # Project vae tokens through moe_gen qkv - vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence) - vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # RoPE on full packed tensor + cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings] + q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0) + k = self.rotary_op(k.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0) - # Reshape to (tokens, heads, head_dim) - text_q = text_q.view(-1, self.num_heads, self.head_dim) - text_k = text_k.view(-1, self.num_kv_heads, self.head_dim) - text_v = text_v.view(-1, self.num_kv_heads, self.head_dim) - vae_q = vae_q.view(-1, self.num_heads, self.head_dim) - vae_k = vae_k.view(-1, self.num_kv_heads, self.head_dim) - vae_v = vae_v.view(-1, self.num_kv_heads, self.head_dim) + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + v = v.to(torch.bfloat16) - # Apply QK norms - text_q = self.q_norm(text_q.to(torch.float32)) - text_k = self.k_norm(text_k.to(torch.float32)) - vae_q = self.q_norm_moe_gen(vae_q.to(torch.float32)) - vae_k = self.k_norm_moe_gen(vae_k.to(torch.float32)) - - # Apply RoPE - need to build per-token cos/sin for text and vae separately - # packed_query_position_embeddings are ordered as the packed sequence - cos_full, sin_full = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings] - - # Extract cos/sin for text and vae positions - text_cos = cos_full[packed_text_indexes] - text_sin = sin_full[packed_text_indexes] - vae_cos = cos_full[packed_vae_token_indexes] - vae_sin = sin_full[packed_vae_token_indexes] - - text_q = self.rotary_op(text_q.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0) - text_k = self.rotary_op(text_k.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0) - vae_q = self.rotary_op(vae_q.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0) - vae_k = self.rotary_op(vae_k.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0) - - text_q = text_q.to(torch.bfloat16) - text_k = text_k.to(torch.bfloat16) - text_v = text_v.to(torch.bfloat16) - vae_q = vae_q.to(torch.bfloat16) - vae_k = vae_k.to(torch.bfloat16) - vae_v = vae_v.to(torch.bfloat16) + # Split into text / vae for SP attention + text_q = q[text_indices] + text_k = k[text_indices] + text_v = v[text_indices] + vae_q = q[vae_indices] + vae_k = k[vae_indices] + vae_v = v[vae_indices] # Build joint K/V: [kv_cache, text_markers] (replicated across SP ranks) if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None: @@ -479,15 +464,19 @@ class PackedAttentionMoT(nn.Module): joint_strategy="front", ), ) - # attn_out: (1, text_len + local_vae_len, H, D) + # attn_out: (1, text_len + vae_len, H, D) text_len = text_q.shape[0] - attn_out = attn_out.squeeze(0) # (text_len + local_vae_len, H, D) + attn_out = attn_out.squeeze(0) text_attn = attn_out[:text_len].reshape(text_len, self.q_size) vae_attn = attn_out[text_len:].reshape(-1, self.q_size) - # Apply output projections - text_out, _ = self.o_proj(text_attn) - vae_out, _ = self.o_proj_moe_gen(vae_attn) + # MoT output projection — construct local packed tensor with local indices + local_packed = torch.cat([text_attn, vae_attn], dim=0) + local_text_idx = torch.arange(text_len, device=local_packed.device) + local_vae_idx = torch.arange(text_len, text_len + vae_attn.shape[0], device=local_packed.device) + local_out, _ = self.o_proj(local_packed, local_text_idx, local_vae_idx) + text_out = local_out[:text_len] + vae_out = local_out[text_len:] # Merge back into packed format total_len = packed_query_sequence.shape[0]