Skip to content

Enable Sequence Parallelism on FP8 Sol-Attn - #35

Merged
lzx1413 merged 1 commit into
Tele-AI:mainfrom
Uxito-Ada:feat/fp8-sol-sequence-parallel-tuning
Aug 25, 2026
Merged

Enable Sequence Parallelism on FP8 Sol-Attn#35
lzx1413 merged 1 commit into
Tele-AI:mainfrom
Uxito-Ada:feat/fp8-sol-sequence-parallel-tuning

Conversation

@Uxito-Ada

@Uxito-Ada Uxito-Ada commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator

Summary

Follow up on #30 by enabling native FP8 Sol-Attn under pure Ulysses sequence parallelism for MiniMax-H3 and Wan

On the matched 4 x H100 MiniMax-H3 workload, TeleFuser FP8 Sol is 2.64x faster end to end, 2.62x faster in denoising, and uses 40.3% less representative peak GPU memory than the deployed LightX2V stack.

Changes

  • Run FP8 Q/K/V after the Ulysses all-to-all, so Q/K block scales and V channel scales cover the exact sequence consumed by each local Sol kernel.
  • Ulysses exchanges BF16 Q/K/V before each rank computes FP8 scales over its full-sequence/local-head shard. Restore after the gathered output.
  • Add unit and 4-GPU NCCL parity coverage for MiniMax-H3 TP2 x SP2 and Wan SP2/SP4.
  • Tune Sol-Attn configs

MiniMax-H3: LightX2V vs TeleFuser

Matched workload: 4 x H100 80GB, TP2 x Ulysses SP2, T2AV 1344x768, 124 frames at 24 fps, 50 configured sampling points (49 denoiser updates), seed 0, one complete warm-up followed by one measured request, feature cache disabled, and MP4 serialization excluded from generation timing.

Prompt:

Steam rises from the ramen while the family talks in the background.

Performance overview

minimax_h3_combo_performance_steps
  • LightX2V uses BF16 Q/K/V inputs dispatch on H100 SM90 to SageAttention2's INT8-QK/FP8-PV kernel with per-thread Q/K quantization, per-channel V quantization, K smoothing, and fp32+fp32 PV accumulation. Its model weights and Linear layers remain BF16.

Generated Output Comparison

All three matched outputs were validated as 1344x768, 124-frame, 24 fps H.264 videos with AAC audio. They generate a coherent family ramen scene, but composition and person placement differ. The saved media is therefore intended for subjective review; performance improvements are not presented as numerical output equivalence.

MiniMax-H3 three-way frame at 2.5 seconds

TeleFuser BF16 + FlashAttention 4

baseline_sp2_tp2_bf16_flash.mp4

TeleFuser W8A8 FP8 + FP8 Sol

optimized_sp2_tp2_fp8_sol_exact.mp4

LightX2V BF16 Linear + SageAttention2 INT8-QK/FP8-PV

lightx2v_matched_50step_768x1344.mp4

TeleFuser FP8 Sol Tuning

Against TeleFuser's matched BF16 baseline, the optimized profile reduces generation time by 33.4%, improves denoising throughput by 53.2%, and reduces total sampled peak memory by 14.4%.

Full-model tuning retained dense_steps=10, dense_layers=2, tau=1.0, and threshold_type=exact. A tau=0.5/diag candidate was slower end to end and did not improve decoded artifact similarity.

H100 microbenchmarks cover 4,096-65,536 global tokens and Ulysses local head counts. Two KV splits remain optimal at and above 16,384 tokens; four splits regress even at 65,536 tokens because workspace and reduction costs outweigh additional KV parallelism. The benchmark and measurements are included under benchmarks/fp8_sol_sequence_parallel/.

Testing

ruff check <PR-scoped files>
# All checks passed

python -m pytest -q \
  tests/unit/models/test_wan_video_sol_attention.py \
  tests/unit/ops/test_fp8_gemm.py \
  tests/unit/ops/test_sol_attention.py \
  tests/unit/pipelines/minimax_h3/test_examples.py \
  tests/unit/pipelines/wan_video/test_optimized_example.py \
  tests/unit/pipelines/wan_video/test_sparse_attention_dispatch.py
# 86 passed

python -m pytest -q \
  tests/integration/test_minimax_h3_distributed.py::test_minimax_h3_fp8_sol_tp2_ulysses2_matches_dense_packed_forward \
  tests/integration/test_wan_fp8_sol_distributed.py
# 3 passed on 4 x H100

Architecture Support

  • Native FP8 Sol kernel: NVIDIA Hopper SM90.
  • Pure Ulysses sequence parallelism: supported.
  • Tensor parallel + Ulysses: supported for MiniMax-H3.
  • Ring and Ulysses-Ring: dense fallback retained.

@Uxito-Ada Uxito-Ada changed the title [PERF] Enable FP8 Sol-Attn with sequence parallelism Enable Sequence Parallelism on FP8 Sol-Attn Aug 24, 2026
@Uxito-Ada
Uxito-Ada requested a review from lzx1413 August 24, 2026 06:33
@Uxito-Ada

Uxito-Ada commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator Author

LightX2V Sol-Attn for Minimax-H3

Futhermore, tested LightX2V's official MiniMax-H3 Sol-Attn example with the corresponding config.

cpu_offload can be disabled, which raises the weights' GPU memory consumption from ~29G to ~79G.

However, the current LightX2V Sol-Attn implementation produces meaningless texture/noise frames, as show below.

image

Env:

Python 3.11.13
Ubuntu 22.04.5
NVIDIA driver 590.48.01
torch                 2.11.0+cu128
torchvision           0.26.0+cu128
torchaudio            2.11.0+cu128
triton                3.6.0
transformers          4.57.3
diffusers             0.37.1
torchada              0.1.82
sol-attn              0.5.0 (editable)
nvidia-cutlass-dsl    4.5.3
sgl-kernel            0.3.21
torchao               0.18.0
cuda-python           13.1.1
cuda-bindings         12.9.7

This indicates an output-correctness bug. Therefore, this PR continues to use LightX2V BF16 + SageAttention2 as the comparison baseline.

@lzx1413

@lzx1413

lzx1413 commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

LGTM

@lzx1413
lzx1413 merged commit 931be41 into Tele-AI:main Aug 25, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants