Skip to content

vllm.models.glm5next.nvidia.specialized_decode.gemm

Skinny BF16 projections for decode steps: the CuTe skinny GEMM that issues its static weight loads before the PDL wait at 1-4 tokens, and Kimi K3's split-K tcgen05 GEMM (cluster DSMEM reduction, evict-first weight TMA) for the DSA and MTP draft projections at 4-64 tokens.

The main model keeps cuBLAS at 8+ tokens on the KDA projections: the skinny GEMM falls behind it there, and the split-K tactics measured slower on the KDA in_proj at B4/B8 (14.4 vs 10.7 us). The DSA projections take the split-K GEMM at 4-64 tokens where their shape has a tactic (policy dsa_proj).

Classes:

  • Projection –

    x @ weight.T for 1, 2 or 4 tokens; split_k_tokens take the split-K

Projection

x @ weight.T for 1, 2 or 4 tokens; split_k_tokens take the split-K GEMM instead where their token tile has a tactic.

Source code in vllm/models/glm5next/nvidia/specialized_decode/gemm.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
class Projection:
    """`x @ weight.T` for 1, 2 or 4 tokens; `split_k_tokens` take the split-K
    GEMM instead where their token tile has a tactic."""

    def __init__(self, weight, split_k_tokens=()):
        self.weight = weight
        n, k = weight.shape
        self.configs = {}
        for m, (block, outputs, unroll) in TACTICS.get((n, k), DEFAULT).items():
            if k % (block * 8) or n % outputs:
                continue
            cfg = SkinnyGemmConfig(m, block, outputs, unroll, 8, k)
            shape_dynamic_skinny_gemm(weight.new_zeros(m, k), weight, cfg, None)
            self.configs[m] = cfg  # compiled before capture
        tactics = SPLIT_K.get((n, k), {})
        policies = {}
        for m in split_k_tokens:
            tile = next((t for t in sorted(tactics) if t >= m), None)
            if tile is not None:
                policies[m] = tactics[tile]
        self.split_k = None
        if policies:
            from .cutedsl.projection import PreparedProjection

            self.split_k = PreparedProjection(weight, policies).prepare()

    def supports(self, tokens: int) -> bool:
        return tokens in self.configs or (
            self.split_k is not None and tokens in self.split_k.kernels
        )

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        m = x.shape[0]
        if self.split_k is not None and m in self.split_k.kernels:
            out = x.new_empty((m, self.weight.shape[0]))
            return self.split_k.run(x, out)
        return shape_dynamic_skinny_gemm(x, self.weight, self.configs[m], None)