vllm.models.glm5next.nvidia.specialized_decode.indexer_query

The indexer's per-token epilogue in one kernel: key LayerNorm, query Hadamard rotation and FP8 quant, head-weight scaling, and the zero padding of the query heads and weights that DeepGEMM's MQA logits need (stock: five kernels).

One warp per query head, four values per lane: the Hadamard stages of strides 1 and 2 stay in a lane, the others pair lanes l and l ^ (stride / 4). The butterflies, the rounding and the power-of-two scale (the stock kernel's accurate log2, read off the exponent bits where that is exact) follow fwht128_quant_fp8 bit for bit.

Functions:

• hadamard_quant –

  One query head, a warp's four values per lane (two BF16x2 words), as

• indexer_query –

  Returns (q_fp8 [M, padded, 128], weights [M, padded], key [M, 128]).

• store_weight –

  w_out[t, head] = value, rounded to BF16 (as .to(torch.bfloat16))

_exponent(x)

ceil(log2(x)) with the stock kernel's accurate log2, for positive x.

From the bits unless the mantissa is barely above a power of two (where the rounded log2 can land on the exponent itself) or x is not normal.

Source code in vllm/models/glm5next/nvidia/specialized_decode/indexer_query.py

@cute.jit
def _exponent(x: Float32):
    """`ceil(log2(x))` with the stock kernel's accurate log2, for positive x.

    From the bits unless the mantissa is barely above a power of two (where the
    rounded log2 can land on the exponent itself) or x is not normal.
    """
    bits = Int32(float_bits(x))
    biased = (bits >> 23) & 255
    mantissa = bits & 0x7FFFFF
    exponent = biased - 127 + Int32(mantissa != 0)
    if (mantissa != 0 and mantissa < 4096) or biased == 0 or biased == 255:
        exponent = Int32(_ceil(cute.math.log2(x, fastmath=False)))
    return exponent

hadamard_quant(got, lane)

One query head, a warp's four values per lane (two BF16x2 words), as fwht128_quant_fp8: (the lane's FP8x4 word, the head's power-of-two scale).

Source code in vllm/models/glm5next/nvidia/specialized_decode/indexer_query.py

@cute.jit
def hadamard_quant(got, lane):
    """One query head, a warp's four values per lane (two BF16x2 words), as
    `fwht128_quant_fp8`: (the lane's FP8x4 word, the head's power-of-two scale).
    """
    x = cute.make_rmem_tensor(cute.make_layout((4,)), Float32)
    for i in cutlass.range_constexpr(2):
        x[2 * i], x[2 * i + 1] = unpack_bf16(got[i])
    # Strides 1 and 2 inside the lane: (a, b) -> (a + b, a - b).
    for stride in (1, 2):
        for i in cutlass.range_constexpr(4):
            if cutlass.const_expr(i % (2 * stride) < stride):
                a, b = x[i], x[i + stride]
                x[i], x[i + stride] = a + b, a - b
    # Strides 4..64 pair lanes l and l ^ (stride / 4); the lower lane holds a.
    for shift in cutlass.range_constexpr(5):
        upper = (lane >> shift) & 1
        for i in cutlass.range_constexpr(4):
            other = cute.arch.shuffle_sync_bfly(x[i], offset=1 << shift)
            if upper == 1:
                x[i] = other - x[i]
            else:
                x[i] = x[i] + other
    amax = Float32(0.0)
    for i in cutlass.range_constexpr(4):
        x[i] = _round_bf16(x[i] * Float32(INV_SQRT_DIM))
        amax = cute.arch.fmax(amax, cute.arch.fmax(x[i], -x[i]))
    amax = cute.arch.fmax(warp_max_abs(amax), Float32(1e-4))
    exponent = _exponent(amax * Float32(1.0 / 448.0))
    q_scale = bits_float(Uint32(exponent + 127) << 23)
    inverse = bits_float(Uint32(127 - exponent) << 23)
    y = [
        cute.arch.fmax(cute.arch.fmin(x[i] * inverse, Float32(448.0)), Float32(-448.0))
        for i in range(4)
    ]
    return Int32(pack_fp8(tuple(y))), q_scale

indexer_query(q, k, weights, norm, scale, padded, weights_dtype=torch.float32)

Returns (q_fp8 [M, padded, 128], weights [M, padded], key [M, 128]).

Parameters:

• q

  (Tensor) –

  [M, heads * 128] BF16 indexer query.

• k

  (Tensor) –

  [M, 128] BF16 indexer key before its LayerNorm.

• weights

  (Tensor) –

  [M, heads] FP32 head weights.

• norm

  (Module) –

  the key LayerNorm module (FP32 weight and bias).

• scale

  (float) –

  the folded softmax and head-count scale.

• padded

  (int) –

  heads after zero padding.

• weights_dtype

  (dtype, default: float32 ) –

  FP32, or BF16 for consumers that round them to BF16.

Source code in vllm/models/glm5next/nvidia/specialized_decode/indexer_query.py

def indexer_query(
    q: torch.Tensor,
    k: torch.Tensor,
    weights: torch.Tensor,
    norm: torch.nn.Module,
    scale: float,
    padded: int,
    weights_dtype: torch.dtype = torch.float32,
):
    """Returns (q_fp8 [M, padded, 128], weights [M, padded], key [M, 128]).

    Args:
        q: [M, heads * 128] BF16 indexer query.
        k: [M, 128] BF16 indexer key before its LayerNorm.
        weights: [M, heads] FP32 head weights.
        norm: the key LayerNorm module (FP32 weight and bias).
        scale: the folded softmax and head-count scale.
        padded: heads after zero padding.
        weights_dtype: FP32, or BF16 for consumers that round them to BF16.

    """
    m, dim = k.shape
    heads = weights.shape[1]
    assert dim == DIM and q.is_contiguous() and k.is_contiguous()
    q_out = torch.empty((m, padded, dim), dtype=torch.float8_e4m3fn, device=q.device)
    w_out = torch.empty((m, padded), dtype=weights_dtype, device=q.device)
    k_out = torch.empty_like(k)
    words = q_out.view(torch.int32).view(m, padded, dim // 4)
    args = (q, k, weights, norm.weight, norm.bias, words, w_out, k_out)
    st = stream(q.device)
    bf16 = weights_dtype == torch.bfloat16
    key = (heads, padded, bf16, q.device.index)
    if key not in _compiled:
        _compiled[key] = cute.compile(
            _QueryKernel(heads, padded, bf16), *(to_cute(t, 16) for t in args),
            Float32(scale), Float32(norm.eps), Int32(m), st,
        )  # fmt: skip
    _compiled[key](
        *(to_cute(t, 16) for t in args),
        Float32(scale), Float32(norm.eps), Int32(m), st,
    )  # fmt: skip
    return q_out, w_out, k_out

store_weight(w_out, t, head, value, padded, bf16)

w_out[t, head] = value, rounded to BF16 (as .to(torch.bfloat16)) when bf16.

Source code in vllm/models/glm5next/nvidia/specialized_decode/indexer_query.py

@cute.jit
def store_weight(w_out, t, head, value, padded: cutlass.Constexpr, bf16):
    """`w_out[t, head] = value`, rounded to BF16 (as `.to(torch.bfloat16)`)
    when `bf16`."""
    if cutlass.const_expr(bf16):
        _store_bf16(w_out.iterator + t * padded + head, value)
    else:
        w_out[t, head] = value