Implementation:Vllm project Vllm SM100 MLA Tile Scheduler

Knowledge Sources	vllm
Domains	Attention, MLA, GPU_Inference
Last Updated	2026-02-08 00:00 GMT

Overview

Implements tile scheduling strategies for SM100 Multi-Latent Attention (MLA) kernels, managing work distribution across GPU thread blocks with both individual and persistent scheduling modes.

Description

This header defines two tile scheduler classes for SM100 MLA attention: Sm100MlaIndividualTileScheduler and Sm100MlaPersistentTileScheduler. The individual scheduler provides a one-to-one mapping between thread blocks and work tiles, while the persistent scheduler enables work-stealing across a fixed grid of thread blocks for better occupancy. Both schedulers compute block coordinates from a three-dimensional problem space (cluster, batch, split_kv).

Usage

These schedulers are instantiated as template parameters of the SM100 MLA FMHA kernel during compilation. The choice between individual and persistent scheduling depends on the problem size and hardware SM count, with the persistent variant preferred for large problems where it can improve GPU utilization.

Code Reference

Source Location

• Repository: vllm
• File: csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp
• Lines: 1-165

Signature

struct Sm100MlaIndividualTileScheduler {
    struct Params { dim3 grid; };

    CUTLASS_DEVICE Sm100MlaIndividualTileScheduler(Params const&);

    template<class ProblemShape, class ClusterShape>
    static Params to_underlying_arguments(
        ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
        ClusterShape const& cluster_shape, int const& split_kv);

    static dim3 get_grid_shape(Params const& params);
    CUTLASS_DEVICE bool is_valid();
    CUTLASS_DEVICE auto get_block_coord();
    CUTLASS_DEVICE Sm100MlaIndividualTileScheduler& operator++();
};

struct Sm100MlaPersistentTileScheduler {
    struct Params {
        int num_blocks;
        FastDivmod divmod_m_block;
        FastDivmod divmod_b;
        FastDivmod divmod_split_kv;
        KernelHardwareInfo hw_info;
    };

    CUTLASS_DEVICE Sm100MlaPersistentTileScheduler(Params const& params);

    template<class ProblemShape, class ClusterShape>
    static Params to_underlying_arguments(
        ProblemShape const& problem_shape, KernelHardwareInfo hw_info,
        ClusterShape const& cluster_shape, int const& split_kv);

    static dim3 get_grid_shape(Params const& params);
    CUTLASS_DEVICE bool is_valid();
    CUTLASS_DEVICE auto get_block_coord();
    CUTLASS_DEVICE Sm100MlaPersistentTileScheduler& operator++();
};

Import

#include "csrc/attention/mla/cutlass_sm100_mla/kernel/sm100_mla_tile_scheduler.hpp"

I/O Contract

Inputs

Name	Type	Required	Description
problem_shape	ProblemShape	Yes	Tuple describing attention dimensions; element at index 3 is the batch size
hw_info	KernelHardwareInfo	Yes	Hardware info containing device_id and sm_count for grid sizing
cluster_shape	ClusterShape	Yes	Cluster tile shape; element at index 0 is the number of M-blocks
split_kv	int	Yes	Maximum number of KV splits for parallel attention computation

Outputs

Name	Type	Description
Params	struct	Scheduler parameters containing grid dimensions and divmod helpers for block coordinate computation
block_coord	cute::tuple	4-element coordinate (m_block, 0, batch_idx, split_kv_idx) identifying the work tile for the current thread block

Usage Examples

// Individual scheduler: one block per tile
using Scheduler = cutlass::fmha::kernel::Sm100MlaIndividualTileScheduler;

auto params = Scheduler::to_underlying_arguments(
    problem_shape, hw_info, cluster_shape, split_kv);
dim3 grid = Scheduler::get_grid_shape(params);

// Inside kernel:
Scheduler scheduler(params);
while (scheduler.is_valid()) {
    auto coord = scheduler.get_block_coord();
    // Process tile at coord...
    ++scheduler;
}

Related Pages

• Environment:Vllm_project_Vllm_CUDA_Runtime