diff --git a/mojo_opset/backends/ttx/kernels/ilu/quant.py b/mojo_opset/backends/ttx/kernels/ilu/quant.py index b8c298fef..cb709d318 100644 --- a/mojo_opset/backends/ttx/kernels/ilu/quant.py +++ b/mojo_opset/backends/ttx/kernels/ilu/quant.py @@ -177,7 +177,8 @@ def _dynamic_quant_row_kernel( abs_max = tl.max(tl.abs(x_scaled), axis=0) qscale = abs_max / 127.0 - inv_qscale = tl.where(qscale > 0.0, 1.0 / qscale, 0.0) + qscale = tl.where(qscale < 1e-6, 1.0, qscale) + inv_qscale = 1.0 / qscale q = x_scaled * inv_qscale # Round-half-away-from-zero, matching the NPU reference kernel. @@ -220,7 +221,8 @@ def _dynamic_quant_row_kernel_blocked( abs_max = tl.maximum(abs_max, cur) qscale = abs_max / 127.0 - inv_qscale = tl.where(qscale > 0.0, 1.0 / qscale, 0.0) + qscale = tl.where(qscale < 1e-6, 1.0, qscale) + inv_qscale = 1.0 / qscale tl.store(qscale_ptr + pid, qscale) for col_off in range(0, n_cols, BLOCK_SIZE_N): diff --git a/mojo_opset/backends/ttx/kernels/ilu/sdpa.py b/mojo_opset/backends/ttx/kernels/ilu/sdpa.py index cc87c0085..aa294b81d 100644 --- a/mojo_opset/backends/ttx/kernels/ilu/sdpa.py +++ b/mojo_opset/backends/ttx/kernels/ilu/sdpa.py @@ -36,7 +36,7 @@ def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=2), ], selected_idx=0, - key=["SEQ", "HEAD_DIM"], + key=["SEQ", "HEAD_DIM", "HAS_MASK"], ) @libentry() @triton.jit @@ -57,6 +57,7 @@ def _sdpa_fav2_fwd_kernel( KV_HEAD_NUM: tl.constexpr, SEQ: tl.constexpr, HEAD_DIM: tl.constexpr, + HAS_MASK: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): @@ -104,14 +105,18 @@ def _sdpa_fav2_fwd_kernel( qk = tl.dot(q_block, k_T_block) qk = qk * sm_scale - attn_mask_ptrs = ( - mask_ptr - + (q_start + offs_m[:, None]) * stride_m0 - + kv_offs[None, :] * stride_m1 - ) mask_valid = ((q_start + offs_m[:, None]) < SEQ) & (kv_offs[None, :] < SEQ) - attn_mask_raw = tl.load(attn_mask_ptrs, mask=mask_valid, other=0) - attn_mask = attn_mask_raw != 0 + if HAS_MASK: + attn_mask_ptrs = ( + mask_ptr + + (q_start + offs_m[:, None]) * stride_m0 + + kv_offs[None, :] * stride_m1 + ) + attn_mask_raw = tl.load(attn_mask_ptrs, mask=mask_valid, other=0) + attn_mask = attn_mask_raw != 0 + else: + # No explicit mask still needs to exclude padded sequence positions. + attn_mask = mask_valid qk = tl.where(attn_mask, qk, float("-inf")) m_ij = tl.maximum(m_i, tl.max(qk, 1)) @@ -160,8 +165,9 @@ def sdpa_infer_impl( assert q.shape[-1] == k.shape[-1] and k.shape[-1] == v.shape[-1] assert q.shape[-2] == k.shape[-2] and k.shape[-2] == v.shape[-2] seq_length = q.shape[-2] - assert mask is not None - assert mask.shape == (seq_length, seq_length) and mask.dtype == torch.bool + has_mask = mask is not None + if has_mask: + assert mask.shape == (seq_length, seq_length) and mask.dtype == torch.bool if not enable_gqa: assert q.shape[1] == k.shape[1] == v.shape[1] @@ -177,7 +183,12 @@ def sdpa_infer_impl( q_c = q.contiguous() k_c = k.contiguous() v_c = v.contiguous() - mask_c = mask.to(torch.int8).contiguous() + if has_mask: + mask_c = mask.to(torch.int8).contiguous() + else: + # The pointer is unused by the HAS_MASK=False specialization. + mask_c = q_c + out = torch.empty_like(q_c) def grid(META): @@ -200,6 +211,7 @@ def grid(META): KV_HEAD_NUM=kv_head_num, SEQ=seq_length, HEAD_DIM=head_dim, + HAS_MASK=has_mask, ) return out.to(q.dtype) diff --git a/mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py b/mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py index 765104969..2ec9b4e91 100644 --- a/mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py +++ b/mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py @@ -40,18 +40,20 @@ def _sdpa_infer_inner( offs_m: tl.constexpr, offs_n: tl.constexpr, # Current stage flag, m and n offset indices SEQ: tl.constexpr, + HAS_MASK: tl.constexpr, fp8_v: tl.constexpr, ): # Iterate over all k, v blocks in the current stage and accumulate the output for start_n in range(0, SEQ, BLOCK_N): # Process BLOCK_N columns at a time start_n = tl.multiple_of(start_n, BLOCK_N) # Align column start position - mask_ptr = ( - mask_base_ptr - + start_m * BLOCK_M * SEQ - + start_n - + tl.arange(0, BLOCK_M)[:, None] * SEQ - + tl.arange(0, BLOCK_N)[None, :] - ) + if HAS_MASK: + mask_ptr = ( + mask_base_ptr + + start_m * BLOCK_M * SEQ + + start_n + + tl.arange(0, BLOCK_M)[:, None] * SEQ + + tl.arange(0, BLOCK_N)[None, :] + ) # -- Compute qk ---- k = tl.load(K_block_ptr) # Modify K @@ -60,11 +62,11 @@ def _sdpa_infer_inner( # NOTE(zhangjihang): tl.where will introduce ub overflow qk = qk * qk_scale - mask = tl.load(mask_ptr) - - # qk += (1 - mask.to(tl.float32)) * (-1e6) - # qk = tl.where(mask, qk, float("-inf")) - qk = tl.where(mask, qk, -1e6) + if HAS_MASK: + mask = tl.load(mask_ptr) + # qk += (1 - mask.to(tl.float32)) * (-1e6) + # qk = tl.where(mask, qk, float("-inf")) + qk = tl.where(mask, qk, -1e6) m_ij = tl.maximum(m_i, tl.max(qk, 1, propagate_nan=tl.PropagateNan.ALL), propagate_nan=tl.PropagateNan.ALL) #m_ij = tl.maximum(m_i, tl.max(qk, 1)) # Scaled max qk = qk - m_ij[:, None] # Stabilize @@ -176,7 +178,7 @@ def get_autotune_config(): triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}), triton.Config({"BLOCK_M": 128, "BLOCK_N": 512}), ], - key=["BSZ", "Q_HEAD_NUM", "SEQ", "HEAD_DIM"], + key=["BSZ", "Q_HEAD_NUM", "SEQ", "HEAD_DIM", "HAS_MASK"], ) @triton.jit def _sdpa_infer_kernel( @@ -208,6 +210,7 @@ def _sdpa_infer_kernel( KV_HEAD_NUM: tl.constexpr, SEQ: tl.constexpr, HEAD_DIM: tl.constexpr, + HAS_MASK: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): @@ -290,6 +293,7 @@ def _sdpa_infer_kernel( offs_m, offs_n, SEQ, + HAS_MASK, V.dtype.element_ty == tl.float8e5, ) @@ -800,8 +804,10 @@ def sdpa_infer_impl( assert head_dim in {64, 128} assert q.shape[-2] == k.shape[-2] and k.shape[-2] == v.shape[-2] seq_length = q.shape[-2] - assert len(mask.shape) == 2 and mask.shape[0] == seq_length and mask.shape[1] == seq_length - assert mask.dtype == torch.bool + has_mask = mask is not None + if has_mask: + assert len(mask.shape) == 2 and mask.shape[0] == seq_length and mask.shape[1] == seq_length + assert mask.dtype == torch.bool if not enable_gqa: assert q.shape[1] == k.shape[1] and q.shape[1] == v.shape[1] @@ -816,6 +822,8 @@ def sdpa_infer_impl( cube_num, vector_num = get_device_properties() num_cores = cube_num M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + # The pointer is unused by the HAS_MASK=False specialization. + mask_arg = mask if has_mask else q # mask = 1 - mask.to(torch.int8) # mask = (1.0 - mask.to(torch.float32)) * (-1e6) @@ -823,7 +831,7 @@ def sdpa_infer_impl( q, k, v, - mask, + mask_arg, M, o, scale, @@ -848,6 +856,7 @@ def sdpa_infer_impl( KV_HEAD_NUM=kv_head_num, SEQ=seq_length, HEAD_DIM=head_dim, + HAS_MASK=has_mask, enable_ubuf_saving=True, enable_hivm_auto_cv_balance=True, multibuffer=True, # 控制开double_buffer diff --git a/mojo_opset/tests/accuracy/operators/test_attention.py b/mojo_opset/tests/accuracy/operators/test_attention.py index 731bad965..d8a277a76 100755 --- a/mojo_opset/tests/accuracy/operators/test_attention.py +++ b/mojo_opset/tests/accuracy/operators/test_attention.py @@ -897,8 +897,11 @@ def generate_diffusion_attn_test_data( @pytest.mark.parametrize( - "bsz, q_head_num, kv_head_num, head_dim, seq_length, block_size", - [(1, 5, 1, 128, 2048, 32,)], + "bsz, q_head_num, kv_head_num, head_dim, seq_length, block_size, mask_kind", + [ + (1, 5, 1, 128, 2048, 32, "diffusion"), + (1, 4, 2, 64, 512, None, "none"), + ], ) @bypass_not_implemented def test_sdpa( @@ -908,17 +911,31 @@ def test_sdpa( head_dim, seq_length, block_size, + mask_kind, ): - query, key, value, blockwise_diffusion_attn_mask, enable_gqa = generate_diffusion_attn_test_data( - bsz, q_head_num, kv_head_num, head_dim, seq_length, block_size - ) - diffusion_attn_ref = MojoSdpa._registry.get("torch")( - scale=1.0 / math.sqrt(query.shape[-1]), enable_gqa=enable_gqa - ) - diffusion_attn = MojoSdpa( - scale=1.0 / math.sqrt(query.shape[-1]), enable_gqa=enable_gqa + if mask_kind == "diffusion": + query, key, value, mask, enable_gqa = generate_diffusion_attn_test_data( + bsz, q_head_num, kv_head_num, head_dim, seq_length, block_size + ) + atol = rtol = 1e-2 + else: + query = torch.randn(bsz, q_head_num, seq_length, head_dim, dtype=torch.bfloat16) * 0.25 + key = torch.randn(bsz, kv_head_num, seq_length, head_dim, dtype=torch.bfloat16) * 0.25 + value = torch.randn_like(key) * 0.25 + mask = None + enable_gqa = q_head_num != kv_head_num + atol, rtol = 6e-2, 8e-2 + + scale = 1.0 / math.sqrt(head_dim) + sdpa = MojoSdpa(scale=scale, enable_gqa=enable_gqa) + sdpa_ref = MojoSdpa._registry.get("torch")(scale=scale, enable_gqa=enable_gqa) + inputs = (query, key, value) if mask is None else (query, key, value, mask) + sdpa.forward_diff_with( + sdpa_ref, + *inputs, + atol=atol, + rtol=rtol, ) - diffusion_attn_ref.forward_diff_with(diffusion_attn, query, key, value, blockwise_diffusion_attn_mask) # =========================================================================== diff --git a/mojo_opset/tests/accuracy/operators/test_gemm.py b/mojo_opset/tests/accuracy/operators/test_gemm.py index edc7c7e54..879994ffa 100644 --- a/mojo_opset/tests/accuracy/operators/test_gemm.py +++ b/mojo_opset/tests/accuracy/operators/test_gemm.py @@ -152,15 +152,19 @@ def test_quant_gemm_parameters_are_registered(): @pytest.mark.parametrize( - "m, k, n", + "m, k, n, output_dtype", [ - (1, 4096, 4096), - (32, 4096, 11008), - (128, 2048, 4096), - (64, 4096, 4096), - ], + (m, k, n, output_dtype) + for m, k, n in [ + (1, 4096, 4096), + (32, 4096, 11008), + (128, 2048, 4096), + (64, 4096, 4096), + ] + for output_dtype in [torch.float16, torch.bfloat16] + ] + + [(65, 257, 129, output_dtype) for output_dtype in [torch.float16, torch.bfloat16, torch.float32]], ) -@pytest.mark.parametrize("output_dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("trans_weight", [False, True]) @bypass_not_implemented @auto_switch_platform() diff --git a/mojo_opset/tests/accuracy/operators/test_position_embedding.py b/mojo_opset/tests/accuracy/operators/test_position_embedding.py index ea0aa9f51..ab0145946 100644 --- a/mojo_opset/tests/accuracy/operators/test_position_embedding.py +++ b/mojo_opset/tests/accuracy/operators/test_position_embedding.py @@ -135,21 +135,14 @@ def test_rotary_embedding(bs, seqlen, rope_dim, mode): @pytest.mark.parametrize( - "bs, seqlen", + "bs, seqlen, dtype, q_heads, k_heads, head_first, head_dim, rope_percentage", [ - (1, 124), - (6, 555), - (2, 2048), - ], -) -@pytest.mark.parametrize( - "dtype, q_heads, k_heads, head_first, head_dim, rope_percentage", - [ - (torch.float16, 32, 8, True, 96, 1.0), - (torch.bfloat16, 8, 2, False, 96, 0.3333333333333333333333), - (torch.float16, 16, 8, True, 128, 1.0), - (torch.bfloat16, 64, 8, False, 88, 1.0), - (torch.float16, 64, 4, True, 128, 0.375), + (1, 124, torch.float16, 32, 8, True, 96, 1.0), + (6, 555, torch.bfloat16, 8, 2, False, 96, 1 / 3), + (2, 2048, torch.float16, 16, 8, True, 128, 1.0), + (1, 124, torch.bfloat16, 64, 8, False, 88, 1.0), + (6, 555, torch.float16, 64, 4, True, 128, 0.375), + (1, 13, torch.bfloat16, 3, 2, False, 64, 0.5), ], ) @pytest.mark.parametrize("mode", ["padding_prefill_pos2d", "padding_prefill_pos3d", "varlen_prefill", "decode"]) diff --git a/mojo_opset/tests/accuracy/operators/test_quantize.py b/mojo_opset/tests/accuracy/operators/test_quantize.py index 1a5bb5fcd..124b98229 100644 --- a/mojo_opset/tests/accuracy/operators/test_quantize.py +++ b/mojo_opset/tests/accuracy/operators/test_quantize.py @@ -61,6 +61,7 @@ (18, 512, [6, 6, 4, 2]), (21, 1024, [2, 5, 1, 7, 6]), (32, 2048, [8, 7, 5, 6, 4, 2]), + (9, 129, [3, 0, 6]), ] dequant_swiglu_quant_cases = [ @@ -197,6 +198,28 @@ def test_dynamic_quant(shape, dtype): quant.forward_diff_with(quant_ref, x, atol=(1, 2e-3), rtol=(0, 2e-3)) +@pytest.mark.parametrize("dtype", dtypes) +@bypass_not_implemented +def test_dynamic_quant_tiny_values_use_unit_scale(dtype): + hidden_size = 129 + x = torch.full((9, hidden_size), 1e-7, dtype=dtype) + inv_smooth_scale = torch.ones(hidden_size, dtype=torch.float32) + quant = load_params( + MojoDynamicQuant(input_size=hidden_size, quant_dtype=torch.int8), + inv_smooth_scale=inv_smooth_scale, + ) + + output, scale = quant(x) + + torch.testing.assert_close(output, torch.zeros_like(output), atol=0, rtol=0) + torch.testing.assert_close( + scale, + torch.ones_like(scale), + atol=0, + rtol=0, + ) + + @pytest.mark.parametrize("tokens, hidden_size, token_count", moe_dynamic_quant_cases) @pytest.mark.parametrize("dtype", dtypes) @bypass_not_implemented