[Perf] [NVIDIA] diffusion: Tune Kitchen INT8 row splitting for SM120

[Tokha233]

Motivation

The Kitchen INT8 row-split default was tuned to avoid expensive Stream-K execution on RTX 4090. On RTX 5090 D v2, splitting H3 QKV matrices into 8192-row calls adds launches and a full output copy and is slower than a single call.

Modifications

Keep a single call by default for SM120, BF16 activations, K <= 8192 and 8192 <= N <= 24832. Keep the existing policy for wider outputs, larger reductions, other dtypes/devices, and whenever either row-split environment variable is explicitly set. Cache the capability by activation device. No weight or arithmetic changes.

The bounds are deliberate: a blanket no-split policy regressed M14850/K7168/N28672 by 11.7%, while K28672 cases were essentially tied. Both remain on the previous path.

Accuracy Test

RTX 5090 D v2 (24GB, SM120), driver 580.126.20, PyTorch 2.13.0+cu130, Triton 3.7.1, stock comfy-kitchen 0.2.36. All 30 tested matrix pairs (15 H3 shapes and 15 boundary/reduction probes) were bit-exact; max_abs=0. BF16 activations and bias, INT8 weights, FP32 row scales, ConvRot256.

14 policy tests passed; all applicable pre-commit hooks passed.

python -m pytest -q python/sglang/multimodal_gen/test/unit/test_kitchen_int8_row_split.py
python -m sglang.multimodal_gen.benchmarks.bench_kitchen_int8_row_split \
  --rows 32700 --input-features 5376 --output-features 16128 --repeats 20

Benchmark & Profiling

Base b906cd3f45bf9d47e92f2d9d7e3d93ad69b06177; candidate d41531d. Matrix sweep uses ABBA, three warmups and ten event-timed repetitions in each arm. Example medians of arm medians:

M/K/N	split8192	single call	latency reduction
8193/5376/16128	3.2457 ms	2.7743 ms	14.52%
14850/5376/16128	5.7207 ms	4.8743 ms	14.79%
32700/5376/16128	12.4758 ms	10.5752 ms	15.23%
32700/5376/8192	6.5016 ms	5.6269 ms	13.45%

No new kernel is introduced; it avoids four GEMMs plus result copies where a single GEMM is more efficient. This is not a universal policy for every Blackwell shape.

Complete SGLang Comfy-integrated Larry v4 INT8 eight-step sampler: 768×512, 124 frames, seed42, Euler/beta, CFG1, sigma shift12/4, torch SDPA, no resident layers. ABBA process order, one warmup and two timed samplers per process. Both arms include the loader correction from #42121 because current main cannot load that serialized checkpoint; this PR does not include the loader change. Both use stock Kitchen 0.2.36.

• Baseline four samples: 33960.216, 33967.192, 33140.687, 33165.233 ms (median 33562.724 ms).
• Candidate four samples: 33365.072, 33370.351, 33421.806, 33426.852 ms (median 33396.079 ms, 0.50% lower).
• All video/audio latent SHA256 values agree across all 12 runs including warmups.
• The reference process spread is larger than the aggregate difference, and its final trial is faster than the candidate. This does not establish stable sampler/serving acceleration. The measured benefit is strongest in the targeted matrix path; no serving-QPS or VAE gain is claimed.
• An existing ComfyUI destructor warning occurs after successful shutdown in all arms; processes exit zero. It is outside timing and is not claimed fixed.

Implementation and verification used Codex assistance. No independent human review is claimed.

Checklist

• Pre-commit formatting/checks.
• Policy regression tests.
• GPU correctness and benchmark, including negative controls.
• Preserve explicit configuration and other hardware defaults.

Raw timings, profiler counts, tests and audit: https://github.com/Tokha233/ComfyUI-H3-SpeedKit/tree/8650e87/evidence/omni-migration-1002

---

CI States

Latest PR Test (Base): ❌ Run #37030652382
Latest PR Test (Extra): ❌ Run #37030651933
Latest PR Test (AMD ROCm 10): ❌ Run #37030652385

[Tokha233]

The GPU workflows stopped at Require run-ci label: the job log reports Missing required label 'run-ci'. Could a maintainer authorize the relevant Diffusion CI?

Local verification on RTX 5090 D v2: 14 policy tests; 30 exact matrix pairs with stock Kitchen 0.2.36; a second run of the checked-in benchmark; and full Larry8 sampler ABBA with all video/audio latent hashes identical. Pre-commit passes. The PR explicitly reports the sampler timing variance and does not claim stable serving throughput improvement.

[Tokha233]

Updated to upstream a74f259 / #38040 in b077a69, resolving the old-file conflict and migrating the policy to convrot_int8_comfy_kitchen.py.

The SM120 policy and explicit environment overrides are unchanged. The benchmark now explicitly constructs the comfy_kitchen backend. The new JIT backend and auto backend selection are untouched; the existing performance figures must not be read as JIT gains.

Validation:

• Policy plus upstream backend-selection suite with CUDA hidden: 32 passed, 8 GPU tests skipped. Changed-file pre-commit passed.
• Reran the checked-in benchmark on RTX 5090 D v2, Torch 2.13.0+cu130, stock Kitchen 0.2.36, M/K/N = 32700/5376/16128. ABBA process-local arm order, 5 warmups + 20 CUDA-event measurements per arm. Arm medians: split 12.460944 / 12.467248 ms; single-call default 10.570080 / 10.544832 ms. Median of arm medians: 12.464096 -> 10.557456 ms, 15.30% lower. Output comparison remains bit-exact.

python -m pytest -q python/sglang/multimodal_gen/test/unit/test_kitchen_int8_row_split.py python/sglang/multimodal_gen/test/unit/test_convrot_int8_backends.py
python -m sglang.multimodal_gen.benchmarks.bench_kitchen_int8_row_split --rows 32700 --input-features 5376 --output-features 16128 --repeats 20

This rerun verifies the migrated Kitchen matrix path, not whole-model or serving throughput. The historical sampler improvement remains 0.50%, within run-to-run variation. GPU CI still requires maintainer run-ci authorization.