[gcanlin]

Purpose

MOSS V2 streaming attention currently produces contiguous [B,H,T,D] output, so the transpose/reshape before the output projection copies data at T>1. Write [B,T,H,D] backing storage directly and return the existing logical shape as a view.

Also add opt-in, offline-selected BF16 Tensor Core FFN GEMMs with GELU / LayerScale+residual epilogues. Dispatch stays inside a custom op so batch/frame dimensions remain symbolic; unmatched shapes retain PyTorch GEMM. The locally retained H200 whitelist selects seven GEMM+GELU shapes that passed speed and waveform checks; no tuning JSON is shipped in this PR. Preserve intermediate BF16 rounding and include configuration contents in the decoder AOT cache key.

• Attention layout is enabled by default. GEMM selection requires MOSS_CODEC_GEMM_CONFIG pointing to a user-supplied JSON file; without it, the original FFN path is used. The tested JSON maps keys "M,3072,768,1" for M in 2/4/8/15/16/30/32 to [16,64,64,1].
• Empty-KV-tile skipping is implemented but disabled by default: sparse histories improve while full windows regress.
• Split-K and other epilogues remain experimental; failed waveform/performance candidates are not in the locally tested whitelist. The local tuner excludes split-K by default.
• Based directly on main 3d952d13; does not include or depend on [Model] Reduce MOSS-TTS codec streaming KV state memory #7332 or the local-depth performance branch. No local experiment logs are added.

Draft for review of the opt-in kernels and validation scope. This PR does not claim an isolated end-to-end throughput or TTFP improvement.

Test Plan

Benchmark runners, validation harnesses, benchmark documentation and tuning JSON under benchmarks/kernels/ are retained locally and intentionally excluded from this PR. Reported measurements were run before removing those files from the PR diff; production code and regression tests are unchanged.

Hardware: one NVIDIA H200 (141 GB), BF16, CUDA graph microbenchmarks. Model waveform validation uses real OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 weights (revision f6e20e543b33d2c252a7ef71bdf8aa71e5ff9169) and fixed random codes, with ten chunks, ring wrap and slot recycling. Generated waveforms are numerical fixtures, not ground-truth speech.

vLLM Version: source 98dff2a81d, version 0.29.0; torch 2.13.0+cu130.

vLLM-Omni Commit: 63bd9ea8, based on main 3d952d13.

In a configured development environment, select an idle GPU and run:

python -m pytest -o addopts='' \
  tests/model_executor/models/moss_tts/test_codec_gemm.py \
  tests/model_executor/models/moss_tts/test_streaming_attention.py -q
pre-commit run --from-ref origin/main --to-ref HEAD

Three warmups; eight calls per CUDA graph, twenty replays per timing sample, median of five samples. GEMM baseline is fullgraph-compiled PyTorch GEMM+epilogue. Whitelist entries require >5% speedup in initial and confirmation timings, plus relative L2 <0.004. Repeated-weight microbenchmarks do not estimate serving throughput.

Test Result

• Standalone PR branch: 14 passed, 15 warnings in 4.97s, including strict dynamic dimensions, graph replay, output strides, SDPA comparison, invalid rows and ring-state handling.
• All applicable pre-commit hooks passed, including mypy, test marks, SPDX, forbidden imports and CUDA-call checks.
• Attention layout alone: 1.24–1.58x for measured T>1 shapes at B=1/4/16; T=1 essentially unchanged. Layout/skip ablations compare bit-exactly with the original kernel.
• GEMM+GELU, N=3072/K=768, microseconds per call:

M	Compiled PyTorch	Selected kernel	Speedup
2	4.464	3.945	1.13x
4	4.820	3.973	1.21x
8	4.796	4.005	1.20x
15	4.851	4.034	1.20x
16	4.795	4.024	1.19x
30	4.969	4.102	1.21x
32	5.049	4.099	1.23x

• Ten-chunk waveform checks passed at B=1 and B=4. Maximum relative RMS: 0.000564 and 0.009101, respectively (numerical gate <0.03). This is not a perceptual quality assessment.
• Perceptual evaluation and isolated serving TTFP/throughput are pending. Existing contributor serving comparisons include other optimizations and are deliberately not attributed to this PR. The local environment emits a vLLM/Omni version warning; validation above records the versions actually used.

AI assistance: Codex assisted with implementation, regression tests, measurements, and PR preparation. This draft is open for human review.