Comfy-Org / Comfy-Org/comfy-kitchen

CUTLASS INT8 dequant selects poor config for tall MiniMax H3 shapes on RTX A6000

Open
#98 2 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
220
Forks
91
Avg merge
1d 7h
Merged PRs (30d)
12

Description

## Summary

The automatic `cutlass_int8_dequant` dispatch appears to select a poor CUTLASS configuration for tall sequence shapes on an NVIDIA RTX A6000 (sm86). The exact MiniMax H3 projection shapes from ComfyUI logs reach the CUTLASS path, but QKV is nearly no faster than BF16 until a faster compiled configuration is selected explicitly.

Environment:

- GPU: NVIDIA RTX A6000, compute capability sm86, 48 GB
- PyTorch: 2.13.0+cu132
- CUDA runtime: 13.2
- comfy-kitchen: 0.2.26
- dtype: bfloat16 output
- ConvRot: true, groupsize 256
- `M = 80661` tokens

## Reproduction

The four production projection shapes are:

```text
M=80661 N=21504 K=5376 # QKV
M=80661 N=5376 K=7168 # attention output
M=80661 N=28672 K=5376 # MLP up
M=80661 N=5376 K=14336 # MLP down + fused SwiGLU
```

ComfyUI logs show the normal path:

```text
[ComfyUI INT8 dispatch] path=int8_quantized_tensor input=(80661, 5376) weight=(21504, 5376) quant=int8_tensorwise convrot=True
[comfy-kitchen INT8] M=80661 N=21504 K=5376 convrot=True input_act=None path=cutlass
```

Production timings for one projection call on the A6000:

```text
QKV BF16: 148.8 ms
QKV INT8 ConvRot: 144.3 ms (1.03x)
```

The extension exposes `benchmark_cutlass_int8_dequant_config`. Sweeping the available configs for the same QKV tensors produced:

```text
config 0: ~84.0 ms kernel-only
config 3: ~98.0 ms
config 2: ~136.5 ms
config 7: ~136.6 ms
config 1: ~137.2 ms
config 12: ~138.0 ms
config 13: ~139.1 ms
```

An end-to-end test using the existing ConvRot activation quantizer plus explicit CUTLASS config 0 measured approximately `89.5 ms`, versus `144.3 ms` through the normal automatic call.

For the same `M`, the fastest configurations were:

```text
QKV N=21504 K=5376: config 0 (~84.0 ms)
Attn out N=5376 K=7168: config 0 (~26.7 ms)
MLP up N=28672 K=5376: config 0 (~111.7 ms)
MLP down N=5376 K=14336: config 13 (~54.5 ms), config 0 ~55.0 ms
```

## Expected behavior

The automatic CUTLASS selection should benchmark or select a configuration based on the actual `(M, N, K, device)` geometry, or use a shape heuristic that handles tall sequence lengths. A fixed/default configuration is especially problematic because `M` is the packed MiniMax sequence length and varies with video duration, spatial resolution, audio length, text, and reference media.

The QKV and MLP-up shapes are architecture-fixed in `N/K`, but `M` changes with the workflow's packed sequence length. The issue is therefore not specific to one literal `M=80661`; the dispatcher should account for tall `M` and wide `N`.

Could you expose the chosen CUTLASS config in diagnostics and add a benchmark/heuristic path for these tall Ampere shapes? The existing benchmark binding made the discrepancy straightforward to verify.

The full benchmark script and logs can be provided if useful.

Contributor guide

Open the contributing guide

Research direction

Start with the existing benchmark_cutlass_int8_dequant_config binding and reproduce the listed QKV and MLP shapes on the RTX A6000. Compare automatic selection with the measured configurations; done means tall-shape dispatch accounts for (M, N, K, device) geometry and exposes the chosen configuration in diagnostics.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.