Grouped tile scheduler drops real output tiles when max_swizzle_size >= 2 and host problem shapes are provided

Description

When a grouped GEMM runs with TileSchedulerArguments::max_swizzle_size >= 2 and host problem shapes are available, the SM90 grouped tile scheduler computes its grid size / early-exit bound without swizzle rounding on the host, but maps linear indices to tiles with swizzle rounding on the device. Whenever any group before the last has swizzle padding (device per-group tile count larger than the cluster-rounded count), all later groups' device-side linear ranges shift upward relative to the host total, and the tail cut drops real output tiles: those CTA positions are never scheduled and the corresponding region of the output matrix is left untouched.

The same scheduler is used underneath PersistentTileSchedulerSm100Group, so Blackwell grouped kernels inherit the behavior.

The two computations

Host side, PersistentTileSchedulerSm90Group::get_tiled_cta_shape_mnl (include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp lines 226-234), rounds each group only to the cluster shape:

auto problem_blocks_m = round_up(ctas_along_m, cute::get<0>(cluster_shape));
auto problem_blocks_n = round_up(ctas_along_n, cute::get<1>(cluster_shape));
total_ctas += problem_blocks_m * problem_blocks_n;

This total becomes both the launched grid and Params::blocks_across_problem_ (include/cutlass/gemm/kernel/tile_scheduler_params.h line 1677).

Device side, the per-group totals include swizzle rounding (sm90_tile_scheduler_group.hpp lines 299-302 and 352-356):

current_group_info_.log_swizzle_size = get_log_swizzle_size(ctas_along_m, ctas_along_n, params_.max_swizzle_size_);
auto problem_blocks_m = round_up(ctas_along_m, (1 << current_group_info_.log_swizzle_size) * params_.cluster_shape_.m());

and the early exit uses the unswizzled host bound (line 441):

if (scheduler_params.pre_processed_problem_shapes && linear_idx >= scheduler_params.blocks_across_problem_) {
    return WorkTileInfo::invalid_work_tile();
}

Repro

Standalone host program replicating both computations exactly (grouped problems 2x2 and 3x3 CTAs, cluster 1x1, max_swizzle_size = 4):

#include <cstdio>
#include <cstdint>
#include <vector>
#include <algorithm>
#include <utility>

static int get_log_swizzle_size(int pm, int pn, int max_sw) {
    int mn = std::min(pm, pn);
    if (max_sw >= 8 && mn >= 6) return 3;
    if (max_sw >= 4 && mn >= 3) return 2;
    if (max_sw >= 2 && mn >= 2) return 1;
    return 0;
}

struct GI { int g; unsigned long long start; unsigned long long total; int ls; };

int main() {
    const int Cm = 1, Cn = 1;
    const int max_swizzle = 4;
    std::vector<std::pair<int,int> > probs;
    probs.push_back(std::make_pair(2,2));
    probs.push_back(std::make_pair(3,3));

    unsigned host_total = 0;
    for (size_t i = 0; i < probs.size(); ++i) {
        int m = probs[i].first, n = probs[i].second;
        host_total += (unsigned)(((m + Cm - 1)/Cm)*Cm) * (((n + Cn - 1)/Cn)*Cn);
    }

    std::vector<GI> gi(probs.size());
    unsigned long long acc = 0;
    for (size_t i = 0; i < probs.size(); ++i) {
        int m = probs[i].first, n = probs[i].second;
        int ls = get_log_swizzle_size(m, n, max_swizzle);
        int mult = (1<<ls);
        unsigned long long pbm = ((m + mult*Cm - 1)/(mult*Cm))*(unsigned long long)(mult*Cm);
        unsigned long long pbn = ((n + mult*Cn - 1)/(mult*Cn))*(unsigned long long)(mult*Cn);
        gi[i] = GI{(int)i, acc, pbm*pbn, ls};
        acc += pbm*pbn;
    }

    printf("host_total=%u dev_total=%llu\n", host_total, acc);

    int dropped = 0;
    for (unsigned long long idx = 0; idx < acc; ++idx) {
        bool rejected_by_early_out = idx >= (unsigned long long)host_total;
        int gsel = -1; GI info;
        for (size_t i = 0; i < gi.size(); ++i) {
            if (idx < gi[i].start + gi[i].total) { gsel = gi[i].g; info = gi[i]; break; }
        }
        if (gsel < 0) continue;
        unsigned long long off = idx - info.start;
        unsigned long long extra  = off >> info.ls;
        unsigned long long sw_off = off & ((1ull<<info.ls)-1);
        int mult = (1<<info.ls);
        int m = probs[gsel].first, n = probs[gsel].second;
        unsigned long long pbm = ((m + mult*Cm - 1)/(mult*Cm))*(unsigned long long)(mult*Cm);
        unsigned long long M = extra % pbm;                       // work_idx_m
        unsigned long long N = (extra / pbm)*mult + sw_off;       // work_idx_n
        bool real = M < (unsigned long long)m && N < (unsigned long long)n;
        if (real && rejected_by_early_out) {
            printf("DROPPED REAL TILE idx=%llu group=%d -> (m=%llu,n=%llu) of %dx%d\n",
                   idx, gsel, M, N, m, n);
            ++dropped;
        }
    }
    printf(dropped ? "FAIL: %d real tiles never scheduled\n" : "OK\n", dropped);
    return 0;
}

Output (g++ -std=c++17 -O2):

host_total=13 dev_total=20
DROPPED REAL TILE idx=13 group=1 -> (m=2,n=1) of 3x3
DROPPED REAL TILE idx=14 group=1 -> (m=2,n=2) of 3x3
FAIL: 2 real tiles never scheduled

With default max_swizzle_size = 1 (log_swizzle_size == 0) the two computations agree and nothing is dropped, which matches why this has not surfaced in stock runs.

Triggering conditions

• Group scheduler (SM90 group kernels directly; SM100/103/120 group paths through PersistentTileSchedulerSm100Group, which forwards to the same SM90 params/device logic)
• Host problem shapes available (pre_processed_problem_shapes == true)
• max_swizzle_size >= 2 with at least one non-final group whose swizzled extent exceeds its cluster-rounded extent

Suggested fix directions

Make the host bound replicate the device per-group computation (including per-group log_swizzle_size rounding) when host shapes are available, or skip swizzle rounding on device when pre_processed_problem_shapes is true. The first keeps the performance benefit and makes grid size consistent with the device map.

Start in include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp, especially get_tiled_cta_shape_mnl and the device-side group total and early-exit logic; also inspect Params::blocks_across_problem_ in include/cutlass/gemm/kernel/tile_scheduler_params.h. Run the standalone C++ reproduction described in the issue. Done means host and device bounds agree and the reproduction reports no dropped real tiles for swizzled grouped shapes.