feat: improve operator accuracy coverage and support mask-free TTX SDPA

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):

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)
 

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,14 +822,16 @@ 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)
     _sdpa_infer_kernel[(num_cores,)](
         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

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)
 
 
 # ===========================================================================

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()

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"])

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