[misc]: extend the comment sweep to every file

2026-11-10 02:07:14 +00:00
parent 74e3061646
commit f496e4b838
155 changed files with 65 additions and 752 deletions
@@ -9,7 +9,6 @@ import triton.language as tl
 from triton.tools.tensor_descriptor import TensorDescriptor

 from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize
-# from attn_qat_infer.api import triton_group_mean


 def is_cuda():
@@ -141,7 +140,6 @@ def _attn_fwd_inner(acc,
    elif STAGE == 2:
        lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
        lo = tl.multiple_of(lo, BLOCK_M)
-    # causal = False
    else:
        lo, hi = 0, N_CTX
    offsetk_y = offset_y + lo  # offset from the start of the current batch-head combination
@@ -149,7 +147,6 @@ def _attn_fwd_inner(acc,
        offsetv_y = offset_y * HEAD_DIM + lo
    else:
        offsetv_y = offset_y + lo
-    # loop over k, v and update accumulator
    for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        kv_valid = (start_n + offs_n) < N_CTX
@@ -191,7 +188,6 @@ def _attn_fwd_inner(acc,
            acc = acc * alpha[:, None]
            if IS_QAT:
                high_prec_acc = high_prec_acc * alpha[:, None]
-        # prepare p and v for the dot
        if dtype == tl.float8e5:
            v = desc_v.load([0, offsetv_y]).T
        else:
@@ -298,7 +294,6 @@ configs = [
    for s in NUM_STAGES_OPTIONS for w in [4, 8]
 ]
 if "PYTEST_VERSION" in os.environ:
-    # Use a single config in testing for reproducibility
    configs = [
        triton.Config(dict(BLOCK_M=128, BLOCK_N=64), num_stages=2, num_warps=4, pre_hook=_host_descriptor_pre_hook),
    ]
@@ -329,8 +324,6 @@ def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape):
        return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape)


-# @triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX_Q", "N_CTX_KV", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"],
-#                  prune_configs_by={'early_config_prune': prune_invalid_configs}, cache_results=True)
@triton.jit
 def _attn_fwd(
    sm_scale,
@@ -398,11 +391,9 @@ def _attn_fwd(
    offset_y_q = off_z * (N_CTX_Q * H) + off_h * N_CTX_Q  # offset for query tensor
    offset_y_kv = off_z * (N_CTX_KV * H) + off_h * N_CTX_KV  # offset for key/value tensors
    qo_offset_y = offset_y_q + start_m * BLOCK_M
-    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = tl.arange(0, BLOCK_N)  # for kv blocks
    q_valid = offs_m < N_CTX_Q
-    # initialize pointer to m and l
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
@@ -411,10 +402,8 @@ def _attn_fwd(
    else:
        # dummy value
        high_prec_acc = acc
-    # load scales
    qk_scale = sm_scale
    qk_scale *= 1.44269504  # 1/log(2)
-    # load q: it will stay in SRAM throughout
    q = desc_q.load([qo_offset_y, 0])  # load from start of q block and start of the entire head dimension
    q = tl.where(q_valid[:, None], q, 0.0)
    # stage 1: off-band
@@ -434,7 +423,6 @@ def _attn_fwd(
                                                           BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV,
                                                           warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P,
                                                           two_level_quant_P, use_global_sf_P, JOIN_QAT_PV)
-    # stage 2: on-band
    if STAGE & 2:
        acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v,
                                                       offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM,
@@ -505,12 +493,10 @@ def _attn_bwd_preprocess(O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, HEAD_
    off_hz = tl.program_id(1)
    off_n = tl.arange(0, HEAD_DIM)
    valid = off_m < N_CTX
-    # load
    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None])
    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :],
                 mask=valid[:, None]).to(tl.float32)
    delta = tl.sum(o * do, axis=1)
-    # write-back
    tl.store(Delta + off_hz * N_CTX + off_m, delta, mask=valid)


@@ -573,7 +559,6 @@ def _attn_bwd_dkdv(
            qk = tl.where(mask, qk, -1.0e6)
        p = tl.math.exp2(qk - m[:, None])
        do = tl.load(do_ptrs, mask=q_valid[:, None])
-        # Compute dV.
        p_quant = p
        if IS_QAT and fake_quant_P:
            p_quant, _ = fake_quantize(src_tensor=p,
@@ -586,7 +571,6 @@ def _attn_bwd_dkdv(
        dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do)
        # D (= delta) is pre-divided by ds_scale.
        Di = tl.load(D + offs_m, mask=q_valid)
-        # Compute dP and dS.
        dp = tl.dot(do, tl.trans(v))
        ds = p * (dp - Di[:, None])
        ds = ds.to(tl.bfloat16)
@@ -659,7 +643,6 @@ def _attn_bwd_dq(
            mask = (offs_m[:, None] >= offs_n[None, :])
            qk = tl.where(mask, qk, -1.0e6)
        p = tl.math.exp2(qk - m)
-        # Compute dP and dS.
        dp = tl.dot(do, tl.trans(v))
        ds = p * (dp - Di[:, None])
        # Compute dQ.
@@ -1037,7 +1020,6 @@ def _attn_bwd(
    if SMOOTH_Q:
        Q_MEAN += adj

-    # load scales
    offs_k = tl.arange(0, HEAD_DIM)

    start_n = pid * BLOCK_N1
@@ -1050,7 +1032,6 @@ def _attn_bwd(
    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)

-    # load K and V: they stay in SRAM throughout the inner loop.
    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])
    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])

@@ -1101,7 +1082,6 @@ def _attn_bwd(
        start_m = 0
        num_steps = (N_CTX + BLOCK_M1 - 1) // BLOCK_M1

-    # Compute dK and dV for non-masked blocks.
    dk, dv = _attn_bwd_dkdv(dk,
                            dv,
                            Q,
@@ -1133,7 +1113,6 @@ def _attn_bwd(
    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dv_ptrs, dv, mask=kv_valid[:, None])

-    # Write back dK.
    dk *= sm_scale
    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dk_ptrs, dk, mask=kv_valid[:, None])
@@ -1185,14 +1164,12 @@ def _attn_bwd(
            warp_specialize=warp_specialize,
        )
        end_n -= num_steps * MASK_BLOCK_N2
-        # stage 2: process KV blocks from 0 to end_n (before diagonal), ceiling division to include remainder
        num_steps = (end_n + BLOCK_N2 - 1) // BLOCK_N2
        start_n = 0
    else:
        # For non-causal, process all KV blocks from 0 to N_CTX (ceiling division to include remainder)
        num_steps = (N_CTX + BLOCK_N2 - 1) // BLOCK_N2
        start_n = 0
-    # stage 2
    dq = _attn_bwd_dq(
        dq,
        q,
@@ -1299,7 +1276,6 @@ class _attention(torch.autograd.Function):
        if IS_QAT:
            high_prec_o = torch.empty_like(q)
        else:
-            # Initialize to a dummy value
            high_prec_o = o
        stage = 3 if causal else 1
        extra_kern_args = {}
@@ -1442,7 +1418,6 @@ class _attention(torch.autograd.Function):
                use_global_sf=use_global_sf_QKV,
            )

-        # Apply pre-hook to set block shapes on tensor descriptors
        _host_descriptor_pre_hook({
            "BLOCK_M": fwd_block_m,
            "BLOCK_N": fwd_block_n,
@@ -1647,7 +1622,6 @@ class _attention(torch.autograd.Function):
                num_stages=dkdv_launch[1],
            )
        elif N_CTX_Q == N_CTX_KV:
-            # Use existing kernel for self-attention (same sequence lengths)
            grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD)
            _attn_bwd[grid](q,
                            arg_k,
@@ -1684,7 +1658,6 @@ class _attention(torch.autograd.Function):
                            num_warps=NUM_WARPS,
                            num_stages=NUM_STAGES)
        else:
-            # Use separate kernels for cross-attention (different sequence lengths)
            grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD)
            _attn_bwd_dq_cross[grid_dq](
                q,
@@ -183,11 +183,9 @@ def _attn_bwd_preprocess(
    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    off_hz = tl.program_id(1)
    off_n = tl.arange(0, HEAD_DIM)
-    # load
    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
    delta = tl.sum(o * do, axis=1)
-    # write-back
    tl.store(Delta + off_hz * N_CTX + off_m, delta)


@@ -255,13 +253,11 @@ def _attn_bwd_dkdv(
        pT = tl.where(mask[:, None], pT, 0.0)

        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
-        # Compute dV.
        ppT = pT
        ppT = ppT.to(tl.bfloat16)
        dv += tl.dot(ppT, do)
        # D (= delta) is pre-divided by ds_scale.
        Di = tl.load(D + offs_m)
-        # Compute dP and dS.
        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
        dsT = pT * (dpT - Di[None, :])
        dsT = dsT.to(tl.bfloat16)
@@ -333,7 +329,6 @@ def _attn_bwd_dq(
        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
        mask = offs_in_block < block_size
        p = tl.where(mask[None, :], p, 0.0)
-        # Compute dP and dS.
        dp = tl.dot(do, vT).to(tl.float32)
        ds = p * (dp - Di[:, None])
        ds = ds.to(tl.bfloat16)
@@ -392,7 +387,6 @@ def _attn_bwd(
    M += off_chz
    D += off_chz

-    # load scales
    offs_k = tl.arange(0, HEAD_DIM)

    start_n = pid * BLOCK_N1
@@ -403,7 +397,6 @@ def _attn_bwd(
    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)

-    # load K and V: they stay in SRAM throughout the inner loop.
    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)

@@ -438,7 +431,6 @@ def _attn_bwd(
    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dv_ptrs, dv)

-    # Write back dK.
    dk *= sm_scale
    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dk_ptrs, dk)
@@ -481,7 +473,6 @@ def _attn_bwd(
        end_n,
        num_steps  #
    )
-    # Write back dQ.
    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
    dq *= sm_scale
    tl.store(dq_ptrs, dq)
@@ -557,7 +548,6 @@ def _attn_bwd_dkdv_kernel(
    start_n = kv_blk * BLOCK_N1
    offs_n = start_n + tl.arange(0, BLOCK_N1)

-    # load K and V: they stay in SRAM throughout the inner loop.
    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)

@@ -165,7 +165,6 @@ def _compute_dequant(
        if dst_dtype == tl.float16:
            dst_scale = dst_scale.to(tl.float16)

-    # Now upcast the tensor.
    intermediate_dtype: tl.constexpr = tl.bfloat16 if dst_dtype == tl.float32 else dst_dtype
    if cuda_capability_geq(10, 0):
        assert is_fp4
@@ -215,7 +214,6 @@ def _compute_dequant(
        # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type
        x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0)
        x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1)
-        # 3) x is zero, do nothing
        dst_tensor = tl.interleave(x0, x1).to(intermediate_dtype, bitcast=True)

    dst_tensor = dst_tensor.to(dst_dtype)
@@ -78,12 +78,10 @@ def _attn_fwd_loop(
    exp_scores = tl.math.exp2(scores - new_m[:, None])
    current_l = tl.sum(exp_scores, axis=1)

-    # Update L <- L * exp(M - M') + L1, M <- M'
    alpha = tl.math.exp2(m - new_m)
    l = l * alpha + current_l
    m = new_m

-    # Update O <- O * exp(M - M') + P @ V
    acc = (acc * alpha[:, None] + tl.dot(exp_scores.to(v.type.element_ty), v))

    return m, l, acc
@@ -286,8 +284,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,

    total_tile_size = tile_t * tile_h * tile_w

-    # BLOCK_Q=128
-    # BLOCK_KV=128
    BLOCK_DIM = head_dim

    output = torch.empty_like(q)
@@ -329,8 +325,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,
                scale=1.0 / (head_dim**0.5),
                has_text=has_text,
                text_q=False,
-                # BLOCK_Q=BLOCK_Q,
-                # BLOCK_KV=BLOCK_KV,
                BLOCK_DIM=BLOCK_DIM,
            )

@@ -363,8 +357,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,
            scale=1.0 / (head_dim**0.5),
            has_text=has_text,
            text_q=True,
-            # BLOCK_Q=BLOCK_Q,
-            # BLOCK_KV=BLOCK_KV,
            BLOCK_DIM=BLOCK_DIM,
        )

@@ -129,21 +129,17 @@ def _rms_norm_fwd_fused(

    x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)

-    # Compute variance
    _var = x * x
    var = tl.sum(_var, axis=1) / N
    rstd = 1 / tl.sqrt(var + eps)

-    # Write mean / rstd
    tl.store(Rstd + rows, rstd)
    rstd = tl.reshape(rstd, (BLOCK_M, 1))

-    # Normalize and apply linear transformation
    w = tl.load(W + cols)
    x_hat = x * rstd
    y = x_hat * w

-    # Write output
    y = y.to(Y.type.element_ty)
    tl.store(y_ptr, y, mask=mask[None, :])

@@ -181,7 +177,6 @@ def rmsnorm(x, w, eps):
    else:
        BLOCK_M = 1

-    # Call the triton kernel
    _rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )](  #
        x,
        y,
@@ -229,20 +224,16 @@ def _layer_norm_param_fwd_fused(

    x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)

-    # Compute mean and Variance
    mean = tl.sum(x, axis=1, keep_dims=True) / N
-    # Compute variance
    _var = (x - mean) * (x - mean)
    var = tl.sum(_var, axis=1, keep_dims=True) / N
    rstd = 1 / tl.sqrt(var + eps)

-    # Write mean / rstd
    _mean = tl.reshape(mean, (BLOCK_M))
    _rstd = tl.reshape(rstd, (BLOCK_M))
    tl.store(Mean + rows, _mean)
    tl.store(Rstd + rows, _rstd)

-    # Normalize and apply linear transformation
    x_hat = (x - mean) * rstd

    w = tl.load(W + cols)
@@ -250,7 +241,6 @@ def _layer_norm_param_fwd_fused(

    x_hat = x_hat * w + b

-    # Write output
    x_hat = x_hat.to(Y.type.element_ty)
    tl.store(y_ptr, x_hat, mask=mask[None, :])

@@ -298,9 +288,6 @@ def layernorm_param(x, w, b, eps):
    return y, mean, rstd


-########################################################
-# Elementwise_affine=False
-########################################################


@triton.jit
@@ -327,23 +314,18 @@ def _layer_norm_noparam_fwd_fused(

    x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)

-    # Compute mean and Variance
    mean = tl.sum(x, axis=1, keep_dims=True) / N
-    # Compute variance
    _var = (x - mean) * (x - mean)
    var = tl.sum(_var, axis=1, keep_dims=True) / N
    rstd = 1 / tl.sqrt(var + eps)

-    # Write mean / rstd
    _mean = tl.reshape(mean, (BLOCK_M))
    _rstd = tl.reshape(rstd, (BLOCK_M))
    tl.store(Mean + rows, _mean)
    tl.store(Rstd + rows, _rstd)

-    # Normalize and apply linear transformation
    x_hat = (x - mean) * rstd

-    # Write output
    x_hat = x_hat.to(Y.type.element_ty)
    tl.store(y_ptr, x_hat, mask=mask[None, :])

@@ -440,7 +422,6 @@ class Int8LinearFunction(torch.autograd.Function):
        grad_bias = None

        if ctx.needs_input_grad[0]:
-            # grad_input = grad_output @ W
            w_float = dequantize_weight(w_q, w_s, grad_output.dtype)
            grad_input = torch.matmul(grad_output, w_float)

@@ -489,34 +470,15 @@ class RMSNormFunction(torch.autograd.Function):
            grad_output_flat = grad_output
            dx = dx

-        # Standard RMSNorm backward
-        # dy = grad_output_flat
-        # x_hat = x * rstd
-        # w * dy
-        # c1 = mean(dy * w * x) * rstd^2
-        # dx = (dy * w - c1 * x) * rstd

-        # More precise:
-        # y = x * rstd * w
-        # dy = grad_output
-        # dw = sum(dy * x * rstd, dim=0)

        grad_w = None
        if ctx.needs_input_grad[1]:  # w
-            # grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0)
-            # More accurate gradient for weight when considering rstd was computed from x
-            # Actually, for Affine part (w), it is just sum(dL/dy * x_hat).
-            # y = x_hat * w
-            # dL/dw = sum(dL/dy * x_hat)
            x_hat = x_flat * rstd
            grad_w = (grad_output_flat * x_hat).sum(dim=0)

        grad_input = None
        if ctx.needs_input_grad[0]:  # x
-            # x_hat = x * rstd
-            # dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat)
-            # but for RMSNorm:
-            # dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula

            # Using PyTorch autograd for reference logic:
            # y = x / sqrt(mean(x^2) + eps) * w
@@ -530,8 +492,6 @@ class RMSNormFunction(torch.autograd.Function):

            # Let's compute:
            dy_w = grad_output_flat * w
-            # term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x)
-            # grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2)

            # Re-derivation:
            # x_hat = x * rstd
@@ -586,10 +546,6 @@ class LayerNormFunction(torch.autograd.Function):
        else:
            w_eff = w

-        # dx calculation
-        # x_hat = (x - mean) * rstd
-        # y = x_hat * w + b
-        # dL/dx_hat = dL/dy * w

        dy = grad_output_flat
        dx_hat = dy * w_eff
@@ -5,7 +5,7 @@ import random
 import time
 import torch
 try:
-    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn import flash_attn_varlen_func
    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
 except ImportError:

@@ -46,7 +46,6 @@ def calc_chunks(cu_seqlen, moba_chunk_size):
    chunk_to_batch[cu_num_chunk[1:-1]] = 1
    chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)

-    # Do not filter out any chunk
    filtered_chunk_indices = torch.arange(num_chunk, device=cu_seqlen.device, dtype=torch.int32)
    num_filtered_chunk = num_chunk

@@ -66,7 +65,7 @@ def _select_threshold_query_head(gate: torch.Tensor, valid_gate_mask: torch.Tens
    eps = 1e-6

    # LSE‐style normalization per <head, query> (across chunks)
-    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    row_min = gate_min_val.amin(dim=0)  # (H, S)
@@ -127,7 +126,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, g
    eps = 1e-6

    # LSE‐style normalization per block (across heads and queries)
-    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    block_max = gate_masked.amax(dim=(1, 2), keepdim=True)  # (C, 1, 1)
@@ -193,7 +192,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
    eps = 1e-6

    # LSE‐style normalization globally across all valid entries
-    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    overall_max = gate_masked.max()  # scalar
@@ -251,7 +250,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,

    # 8) create the mask in the original flat shape
    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool)  # (C*H*S,)
-    if selected_other_indices.numel() > 0:  # Check if any 'other' indices were selected
+    if selected_other_indices.numel() > 0:
        others_mask_flat[selected_other_indices] = True
    others_mask = others_mask_flat.view(C, H, S)  # (C, H, S)

@@ -349,8 +348,6 @@ class MixedAttention(torch.autograd.Function):
        ctx.moba_chunk_size = moba_chunk_size
        ctx.softmax_scale = softmax_scale = q.shape[-1]**(-0.5)

-        # Non-causal self-attention branch
-        # return out, softmax_lse, S_dmask, rng_state
        self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
            q=q,
            k=k,
@@ -532,12 +529,10 @@ def moba_attn_varlen(
    kv = torch.stack((k, v), dim=1)
    seqlen, num_head, head_dim = q.shape

-    # Compute chunk boundaries.
    cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(cu_seqlens, moba_chunk_size)

    self_attn_cu_seqlen = cu_chunk

-    # Update top-k selection to include the self chunk.
    moba_topk = min(moba_topk, num_filtered_chunk)

    # --- Build filtered KV from chunks ---
@@ -553,16 +548,13 @@ def moba_attn_varlen(
    gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
    gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)

-    # Compute key_gate_weight over valid tokens.
    key_values = gathered[:, :, 0].float()  # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
    valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
    key_sum = (key_values * valid_mask_exp).sum(dim=1)
    divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
    key_gate_weight = key_sum / divisor  # [num_filtered_chunk, num_head, head_dim]

-    # Compute gate logits between key_gate_weight and queries.
    q_float = q.float()
-    # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float)  # [num_filtered_chunk, num_head, seqlen]
    gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)

    # Amplify the diagonal (self chunk) contributions.
@@ -650,7 +642,6 @@ def moba_attn_varlen(
    moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
    moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)

-    # Build cumulative sequence lengths for the selected queries.
    moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
    q_zero_mask = moba_seqlen_q == 0
    valid_expert_mask = ~q_zero_mask
@@ -814,7 +805,6 @@ def test_attn_varlen_moba_speed(batch,
                                simsum_threshold=0.25,
                                threshold_type='query_head'):
    """Speed test comparing flash_attn vs moba_attention"""
-    # Get data
    q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
    print(
        f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}"