daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Record exact FC2 cuBLASLt schedule

Browse source
This commit is contained in:
Daniel Maddern 2026-08-25 22:32:48 +07:00
parent ea14b98544
commit 0298ccd1cb
20 changed files with 12886 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
CURRENT_STATE.md
View file

@ -1,6 +1,6 @@
# H3 Runtime Current State # H3 Runtime Current State
Status date: 2026-08-24 Status date: 2026-08-25
This document is the canonical snapshot of implemented scope and remaining work. This document is the canonical snapshot of implemented scope and remaining work.
Historical handoffs in `PLAN.md` and `PARITY.md` may describe older states. Historical handoffs in `PLAN.md` and `PARITY.md` may describe older states.
@ -160,6 +160,18 @@ video and audio tensors. Spark enables the path with
`SAGE2_BLACKWELL_DESIGN.md`, and `SAGE2_BLACKWELL_DESIGN.md`, and
`benchmarks/gb10-post-optimization-profile-summary.json`. `benchmarks/gb10-post-optimization-profile-summary.json`.
The isolated FC2 cuBLASLt scheduling study is complete. The production
heuristic's `_stream_k` kernel requests the same `25.664 GB` of operands as the
retained public split-K-1 schedule, but its L2 hit rate is only `53.32%` versus
`91.10%`; it incurs `9.853 GB` more L2 read misses and spends heavily in
synchronization polling. Algorithm 70, tile 20, stages 37, split-K 1 is
byte-exact with zero workspace. It improves complete blocks 0, 24, and 49 by
`8.16-8.88%`, the two-step trajectory by `7.50%`, and the canonical 12-step
trajectory from `278.201 s` to `255.371 s` (`8.21%`) with exact video and audio
latents. This remains a research-retained integration candidate: production
dispatch and configuration are unchanged. See
`research/fc2_nvfp4_scheduling/RESULTS.md`.
The Spark hot runtime was rebuilt and recreated with image The Spark hot runtime was rebuilt and recreated with image
`sha256:5f879c43374bcedb89745971d7d95d82afc8fcf9c41d30f11b257c2863b9fe28`. `sha256:5f879c43374bcedb89745971d7d95d82afc8fcf9c41d30f11b257c2863b9fe28`.
Health and startup warmup pass with modulation fusion enabled. A resident real Health and startup warmup pass with modulation fusion enabled. A resident real

21
PERFORMANCE_ROADMAP.md
View file

@ -200,7 +200,8 @@ another's result.
### Phase 3: NVFP4 Fused Projection Prototype ### Phase 3: NVFP4 Fused Projection Prototype
Status: active next component on GB10/SM121. Status: active on GB10/SM121; isolated FC2 library scheduling is complete and
awaits production integration.
The first profile also exposed and fixed a native activation-packer layout bug: The first profile also exposed and fixed a native activation-packer layout bug:
the previous width-specific block-scale swizzle failed at the attention output's the previous width-specific block-scale swizzle failed at the attention output's
@ -383,6 +384,24 @@ off-chip traffic bottleneck, the NVFP4 GEMMs. Do not select a new kernel from
the old profile. See the old profile. See
`benchmarks/gb10-fully-fused-fresh-nsight-summary.json`. `benchmarks/gb10-fully-fused-fresh-nsight-summary.json`.
The FC2 follow-up resolves that projection's `16.1x` off-chip amplification.
The production heuristic launches an undocumented-sentinel `_stream_k` kernel;
the retained documented configuration is cuBLASLt algorithm 70, tile 20,
stages 37, public split-K 1, reduction scheme 0, and zero workspace. Both
schedules request the same `25.664 GB` of operands, but split-K 1 raises L2 hit
rate from `53.32%` to `91.10%`, removes `9.853 GB` of L2 read misses, and raises
tensor-pipe activity from `25.01%` to `82.42%`. There is no material global
partial-accumulator or output-reduction traffic; the baseline instead spends
heavily in Stream-K synchronization polling and loses traversal locality.
The candidate is byte-exact and improves FC2 p50 from `52.521 ms` to
`15.636 ms`. Complete blocks 0, 24, and 49 improve by `8.16-8.88%`; two-step
and canonical 12-step trajectories improve by `7.50%` and `8.21%`, with exact
video and audio latents. The custom persistent-kernel branch is therefore
closed. Production integration remains separate work, so current dispatch is
unchanged. See `research/fc2_nvfp4_scheduling/RESULTS.md` and
`benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json`.
The follow-up real block-24 Sage2 decomposition now selects the next exact The follow-up real block-24 Sage2 decomposition now selects the next exact
kernel experiment. Manual preparation plus the existing prequantized mainloop kernel experiment. Manual preparation plus the existing prequantized mainloop
is byte-exact against public SageAttention 2.2.0. Uninstrumented median phase is byte-exact against public SageAttention 2.2.0. Uninstrumented median phase

4
benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv Normal file
View file

File diff suppressed because one or more lines are too long

BIN
benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep Normal file
View file

Binary file not shown.

1733
benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

2014
benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

2133
benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

3
benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv Normal file
View file

File diff suppressed because one or more lines are too long

BIN
benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep Normal file
View file

Binary file not shown.

1741
benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

1769
benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

1769
benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json Normal file
View file

File diff suppressed because it is too large Load diff

188
research/artifact_manifest.json
View file

@ -2,39 +2,199 @@
"metadata": { "metadata": {
"algorithm": "SHA-256", "algorithm": "SHA-256",
"generated_date": "2026-08-25", "generated_date": "2026-08-25",
"scope_note": "Local selected scope is the established 167-file set. Spark records are the complete reproducible current top-level benchmark output set; the audit-retained 279-file/665950155-byte aggregate cannot be reconstructed because its path list was not retained.", "scope_note": "Local selected scope is the established 167-file set plus 10 retained FC2 scheduling artifacts. Spark records are the complete reproducible current top-level benchmark output set; the audit-retained 279-file/665950155-byte aggregate cannot be reconstructed because its path list was not retained.",
"summary": { "summary": {
"local": { "local": {
"record_count": 167, "record_count": 177,
"size_bytes": 210388688, "size_bytes": 230157548,
"expected_record_count": 167, "expected_record_count": 177,
"expected_size_bytes": 210388688, "expected_size_bytes": 230157548,
"reconciled": true "reconciled": true
}, },
"spark": { "spark": {
"record_count": 280, "record_count": 290,
"size_bytes": 547708583, "size_bytes": 567477443,
"expected_record_count": 279, "expected_record_count": 279,
"expected_size_bytes": 665950155, "expected_size_bytes": 665950155,
"reconciled": false, "reconciled": false,
"record_count_delta": 1, "record_count_delta": 11,
"size_bytes_delta": -118241572 "size_bytes_delta": -98472712
}, },
"total": { "total": {
"record_count": 447, "record_count": 467,
"size_bytes": 758097271 "size_bytes": 797634991
} }
}, },
"json_reconciliation": { "json_reconciliation": {
"identical": 106, "identical": 112,
"mismatches": 2, "mismatches": 2,
"local_only": 40, "local_only": 40,
"spark_only": 107, "spark_only": 107,
"local_total": 148, "local_total": 154,
"spark_total": 215 "spark_total": 221
} }
}, },
"artifacts": [ "artifacts": [
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv",
"size_bytes": 164324,
"sha256": "e5bfd39f1512737c0852975de319b502e2ce20677e2a51ad0ce50656923634a5",
"artifact_class": "nsight_csv_export",
"location": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep",
"size_bytes": 13364977,
"sha256": "975eea7110a19b472ef8cd29ec1c72f629b0f97e224080acc50a43b371bc290d",
"artifact_class": "nsight_compute_report",
"location": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json",
"size_bytes": 51924,
"sha256": "082e195a34a0a4738cd515dcc7db69e59d8bef22afcb49e73a91d21cb0ccaff0",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json",
"size_bytes": 61451,
"sha256": "d1051463765052c83494f417b40b9fe9eb389eb4ebf213ff7ff491991d569d8c",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json",
"size_bytes": 64494,
"sha256": "c762e9394f6ab6dd1edfd13097a7fa242caa4540fb3d71c3b1035b3edbcd3261",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv",
"size_bytes": 148497,
"sha256": "2a9da47628571fd430a8c0ba4fce3547d8e12dec3b443b101a3a8baa4bdf8e8f",
"artifact_class": "nsight_csv_export",
"location": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep",
"size_bytes": 5755084,
"sha256": "b368a14dcd96f859dbad9e7a8fa10bc33bf6373dbea8972b6685a2d83b68dfae",
"artifact_class": "nsight_compute_report",
"location": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json",
"size_bytes": 52089,
"sha256": "7b2282b6381e90d07d7175fe3a8f1aa03c7d85a3c6d25cab2350abf6559800d8",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json",
"size_bytes": 53013,
"sha256": "075c3aecbe6f77cc28f07a8f656de044e437c94a9fc71068c0ac9e8f3082972c",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json"
},
{
"scope": "local",
"path": "benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json",
"size_bytes": 53007,
"sha256": "a9cc0b229dcc7e37705d94705b1e607092d9aee7e3873787eff539074f46228f",
"artifact_class": "benchmark_json",
"location": "benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv",
"size_bytes": 164324,
"sha256": "e5bfd39f1512737c0852975de319b502e2ce20677e2a51ad0ce50656923634a5",
"artifact_class": "nsight_csv_export",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep",
"size_bytes": 13364977,
"sha256": "975eea7110a19b472ef8cd29ec1c72f629b0f97e224080acc50a43b371bc290d",
"artifact_class": "nsight_compute_report",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json",
"size_bytes": 51924,
"sha256": "082e195a34a0a4738cd515dcc7db69e59d8bef22afcb49e73a91d21cb0ccaff0",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json",
"size_bytes": 61451,
"sha256": "d1051463765052c83494f417b40b9fe9eb389eb4ebf213ff7ff491991d569d8c",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json",
"size_bytes": 64494,
"sha256": "c762e9394f6ab6dd1edfd13097a7fa242caa4540fb3d71c3b1035b3edbcd3261",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv",
"size_bytes": 148497,
"sha256": "2a9da47628571fd430a8c0ba4fce3547d8e12dec3b443b101a3a8baa4bdf8e8f",
"artifact_class": "nsight_csv_export",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep",
"size_bytes": 5755084,
"sha256": "b368a14dcd96f859dbad9e7a8fa10bc33bf6373dbea8972b6685a2d83b68dfae",
"artifact_class": "nsight_compute_report",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json",
"size_bytes": 52089,
"sha256": "7b2282b6381e90d07d7175fe3a8f1aa03c7d85a3c6d25cab2350abf6559800d8",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json",
"size_bytes": 53013,
"sha256": "075c3aecbe6f77cc28f07a8f656de044e437c94a9fc71068c0ac9e8f3082972c",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json"
},
{
"scope": "spark",
"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json",
"size_bytes": 53007,
"sha256": "a9cc0b229dcc7e37705d94705b1e607092d9aee7e3873787eff539074f46228f",
"artifact_class": "benchmark_json",
"location": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json"
},
{ {
"scope": "local", "scope": "local",
"path": "benchmarks/audio-dialogue-format-sweep-sage2-seeds440420-440429.json", "path": "benchmarks/audio-dialogue-format-sweep-sage2-seeds440420-440429.json",

22
research/experiment_registry.json
View file

@ -500,6 +500,28 @@
"reproducer_commands": [], "timestamp": null, "evidence_missing": ["No accepted alternative FC2 reduction implementation"], "reproducer_commands": [], "timestamp": null, "evidence_missing": ["No accepted alternative FC2 reduction implementation"],
"production_behavior": "FC2 retains the Comfy/CUBLAS fallback.", "source_recovery": "P0 artifacts and fallback policy are retained in the roadmap." "production_behavior": "FC2 retains the Comfy/CUBLAS fallback.", "source_recovery": "P0 artifacts and fallback policy are retained in the roadmap."
}, },
{
"id": "fc2-cublaslt-splitk1-schedule",
"name": "FC2 cuBLASLt public split-K-1 schedule",
"family": "nvfp4-library-scheduling",
"status": "research_retained",
"hypothesis": "A documented cuBLASLt schedule can preserve the exact FC2 reduction result while avoiding the production heuristic's traffic and synchronization regression.",
"implementation_strategy": "Reproduce the exact Comfy Kitchen descriptors in an isolated extension, enumerate checked cuBLASLt algorithms, and compare one selected public split-K-1 schedule against the accepted FC2 path.",
"source_locations": ["research/fc2_nvfp4_scheduling/README.md", "research/fc2_nvfp4_scheduling/RESULTS.md", "research/fc2_nvfp4_scheduling/fc2_nvfp4_lt.cpp", "tools/benchmark_fc2_nvfp4_algorithms.py"],
"active_source_location": "research/fc2_nvfp4_scheduling/fc2_nvfp4_lt.cpp",
"commit_hash": null,
"benchmark_artifacts": [{"path": "benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json", "exists": true}],
"profiler_artifacts": [{"path": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv", "exists": true}, {"path": "benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv", "exists": true}],
"environment": {"gpu": "NVIDIA GB10", "cuda": "13", "driver": null, "pytorch": null, "triton": null, "container_image": "sha256:1d340e14cb6fc45ccfdbe63dde8db2a2b3aea94b493702c8a08e1f8d5b4f7b83", "commit_context": "isolated extension; production source and dispatch unchanged"},
"metrics": {"fc2_p50_ms_baseline": 53.618, "fc2_p50_ms_candidate": 15.636, "block_improvement_pct": [8.16, 8.63, 8.88], "two_step_improvement_pct": 7.50, "canonical_12_step_improvement_pct": 8.21, "trajectory_measurement_note": "single baseline-then-candidate paired runs", "l2_hit_rate_pct_baseline": 53.32, "l2_hit_rate_pct_candidate": 91.10, "l2_read_miss_reduction_bytes": 9853094784},
"correctness_evidence": ["Selected FC2 output is byte-exact.", "Blocks 0, 24, and 49 are byte-exact against paired baseline and retained traversal.", "Two-step and canonical 12-step video and audio latents are byte-exact."],
"decision_rationale": "The library schedule passed all exactness and performance gates, explaining the prior traffic amplification without requiring a custom kernel. It is retained pending explicit production integration and deployment validation.",
"reproducer_commands": ["python tools/benchmark_fc2_nvfp4_algorithms.py --mode sweep --rounds 20 --workspace-bytes 67108864", "python tools/benchmark_fc2_nvfp4_algorithms.py --mode block-gate --candidate research/fc2_nvfp4_scheduling/candidate_splitk1.json --rounds 20 --workspace-bytes 0"],
"timestamp": "2026-08-25",
"evidence_missing": ["Production Nvfp4Linear integration", "Production deployment smoke", "Portable validation outside GB10/SM121"],
"production_behavior": "Not dispatched. Production still uses Comfy Kitchen's heuristic-selected FC2 path.",
"source_recovery": "The direct cuBLASLt extension, benchmark harness, selected candidate, raw NCU reports, and full gate artifacts are retained in this checkout."
},
{ {
"id": "layout-direct-temporal-output", "id": "layout-direct-temporal-output",
"name": "Direct-to-temporal output", "name": "Direct-to-temporal output",

206
research/fc2_nvfp4_scheduling/README.md Normal file
View file

@ -0,0 +1,206 @@
# FC2 NVFP4 Library Scheduling
This directory contains the first isolated library-scheduling study for H3 FC2.
It does not register a PyTorch operator, alter `Nvfp4Linear`, or participate in
production dispatch. The extension is loaded only by
`tools/benchmark_fc2_nvfp4_algorithms.py`.
The completed measurements and decision are in `RESULTS.md`. Algorithm 70 with
public split-K 1 is byte-exact and passes the block, two-step, and canonical
12-step gates, but remains research-only until production integration and
deployment validation are performed.
## Exact Operation
The extension reproduces Comfy Kitchen 0.2.31's
`cublas_gemm_nvfp4.cu` descriptors for row-major packed activation `[M,K]`
times packed weight `[N,K]` to BF16 `[M,N]`:
- cuBLASLt sees column-major `weight.T @ activation`, so Lt `m=N`, `n=M`, and
`k=K`; the output storage is the row-major `[M,N]` view.
- A and B are `CUDA_R_4F_E2M1`, with
`CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3` block scales.
- Compute, scalar scale type, alpha, and beta are FP32. Alpha and zero beta are
device pointers.
- Output is BF16, epilogue is default, and there is no bias.
- The activation producer is the deployed, accepted
`vortex_native_quantize_swiglu_nvfp4`. Alpha is its FP32 tensor scale times
FC2's FP32 weight tensor scale.
The canonical 37,810-row input pads to 37,824 rows in the accepted producer.
The report distinguishes logical `[M,N,K]` from descriptor/padded dimensions;
comparison slices back to the logical output exactly as Comfy Kitchen does.
## Current Baseline
The accepted production path fuses exact BF16 SwiGLU into NVFP4 production and
retains Comfy Kitchen's cuBLASLt FC2 GEMM. Existing validated metadata reports:
- Producer tensor scale, QDATA, and SFA are byte-identical at blocks 0, 24, and
49.
- Canonical block FC2 input is logically `37,810 x 14,336` BF16 before packing.
- Fully fused block-24 profiling attributes 26.87% of kernel time to all four
NVFP4 GEMMs; this is not claimed as an FC2-only percentage.
- Spark's accepted image is
`sha256:1d340e14cb6fc45ccfdbe63dde8db2a2b3aea94b493702c8a08e1f8d5b4f7b83`
and enables `H3_NVFP4_SCALE_BACKEND=vortex`,
`H3_FUSED_ELEMENTWISE=1`, `H3_NVFP4_MODULATE_FUSION=1`, and
`H3_NVFP4_SWIGLU_FUSION=1`.
Every run also records a fresh profiler-derived baseline kernel list. Kernel
names and times in that list are run metadata, not hard-coded claims.
## Commands
Run inside image `1d340e14cb6f...` from the Spark checkout, without starting the
hot service. Preserve its production environment switches, for example:
```bash
export H3_NVFP4_SCALE_BACKEND=vortex H3_NVFP4_SCALE_VERSION=1
export H3_FUSED_ELEMENTWISE=1 H3_NVFP4_MODULATE_FUSION=1
export H3_NVFP4_SWIGLU_FUSION=1 H3_SAGE_QKV_LAYOUT=strided_nhd
```
Compile only:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode compile --verbose-build \
--build-directory /tmp/fc2-nvfp4-build
```
Default characterization enumerates workspace budgets 0, 4, 8, 16, 32, and
64 MiB, checks a bounded set of public split-K configurations, and times one
valid probe with alternating AB/BA order:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode characterize \
--build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-characterize.json
```
The same workload can come from an existing capture:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode characterize \
--capture /artifacts/capture --build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-capture.json
```
A bounded timing sweep defaults to eight candidates, not a full combinatorial
gate:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode sweep \
--candidate-limit 8 --rounds 6 --workspace-bytes 33554432 \
--build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-sweep.json
```
Re-run one selected candidate. `--candidate` accepts an inline JSON object, a
JSON file containing one object, or `enumerate:<budget-index>:<result-index>`:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode selected \
--candidate '{"algorithm_id":23,"tile_id":42,"stages_id":35,"split_k":1,"reduction_scheme":0}' \
--workspace-bytes 33554432 --build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-selected.json
```
NCU capture uses `cudaProfilerStart/Stop`; only the selected FC2 cuBLASLt call
is inside the range:
```bash
ncu --target-processes all --profile-from-start off --set full \
--export /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-selected \
python tools/benchmark_fc2_nvfp4_algorithms.py --mode profile \
--candidate /tmp/fc2-candidate.json --workspace-bytes 33554432 \
--build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-ncu.json
```
Complete block gate for blocks 0, 24, and 49:
```bash
python tools/benchmark_fc2_nvfp4_algorithms.py --mode block-gate \
--candidate /tmp/fc2-candidate.json --rounds 6 \
--workspace-bytes 33554432 --build-directory /tmp/fc2-nvfp4-build \
--output /output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-block-gate.json
```
Block-gate monkeypatches only `block.mlp.fc2.forward_swiglu`. It still invokes
the accepted SwiGLU producer on the actual FC1 result. Norm, FC1, modulation,
gate, residual, attention, and all other code remain the accepted path.
## Requested NCU Traffic Counters
Use `ncu --query-metrics` on the installed NCU version before replacing `--set
full` with an explicit list. Request the available equivalents of:
```text
gpu__time_duration.sum
dram__bytes_read.sum
dram__bytes_write.sum
lts__t_bytes.sum
lts__t_sector_hit_rate.pct
sm__throughput.avg.pct_of_peak_sustained_elapsed
smsp__inst_executed.sum
smsp__pipe_tensor_cycles_active.avg.pct_of_peak_sustained_active
launch__registers_per_thread
launch__shared_mem_per_block_allocated
```
Metric spelling and availability vary with NCU and GB10. Missing counters must
be reported as unavailable rather than silently substituted. Also retain
kernel duration, grid/cluster dimensions, achieved occupancy, waves, required
workspace, and the extension's tile/stage/inner/cluster configuration.
## Validation Ladder
1. Compile against the image's CUDA 13 headers and print `build_info()`.
2. Characterize heuristics and verify every executed candidate passes
`cublasLtMatmulAlgoCheck` within caller-supplied workspace.
3. Compare direct-probe candidate FC2 against a cloned
`fc2.forward_swiglu(gate_up)` output byte-for-byte.
4. Alternate baseline/candidate AB and BA each round; record p50, p95, dense
TFLOP/s, producer-plus-FC2 time, workspace, and output SHA-256.
5. NCU one FC2 call and inspect traffic/resource counters.
6. Gate complete blocks 0, 24, and 49 with only FC2 GEMM replaced. Require
BF16 exactness against both the paired baseline and retained traversal.
7. Only after those gates should a separate, explicitly authorized trajectory
experiment be considered.
That trajectory authorization was granted for the retained candidate. Both the
two-step and canonical 12-step gates passed with byte-exact video and audio
latents; see `RESULTS.md` and the linked benchmark artifacts.
## Prohibited Experiments
- Do not alter production source, configuration, dispatch, or existing
research files.
- Do not expose this extension through `Nvfp4Linear` or any production operator.
- Do not start or perturb the hot service.
- Do not add bias, accumulation, host scalars, another output dtype, a different
producer, or approximate validation to this study.
- Do not claim a checkpoint hash unless `--checkpoint-sha256` actually computes
it.
- Do not label undocumented behavior Stream-K. CUDA 13 exposes no documented
public `cublasLtMatmulAlgoConfig` attribute that directly selects Stream-K;
the report states that limitation explicitly.
- Do not promote a direct-GEMM result without the complete block gate and later
separately authorized trajectory validation.
## Limitations
- Heuristics are library, driver, GPU, shape, and workspace specific.
- Explicit split-K is attempted only through documented `AlgoInit`,
`ConfigSet`, and `AlgoCheck`. Unsupported factors/reduction schemes are
recorded; they are not emulated.
- Negative `SPLITK_NUM` values returned by the heuristic are retained as raw
signed library sentinels. They are not interpreted or labeled as Stream-K.
- Characterization deliberately bounds explicit checks and timing candidates.
Increase limits consciously because each full canonical FC2 call is costly.
- Caller-owned output and workspace are reused during timing. The accepted
producer still owns its quantized activation allocations.
- A capture supplies packed-denoiser inputs, not pre-captured FC2 operands; the
harness traverses the loaded H3 model once to construct exact current
boundaries.

158
research/fc2_nvfp4_scheduling/RESULTS.md Normal file
View file

@ -0,0 +1,158 @@
# FC2 NVFP4 Scheduling Results
Status: research retained; validated as a canonical-shape production-integration
candidate, but not wired into production dispatch.
## Decision
The canonical GB10 FC2 slowdown is caused by the cuBLASLt heuristic selecting a
`_stream_k` kernel with poor traversal locality and heavy synchronization
polling. A documented public split-K-1 configuration is byte-exact and removes
the regression. Because this library candidate passed the complete validation
ladder, the custom persistent-kernel branch is closed.
The retained candidate is:
| Field | Value |
| --- | ---: |
| Algorithm ID | `70` |
| Tile ID | `20` |
| Stages ID | `37` |
| Split-K | `1` |
| Reduction scheme | `0` |
| Custom option / CTA swizzle | `0 / 0` |
| Required workspace | `0 bytes` |
| Supplied workspace | `64 MiB` sweep; `0 bytes` block/trajectory gates |
The baseline heuristic's negative split-K value is an undocumented library
sentinel. CUDA 13 has no documented public attribute that directly selects
Stream-K, so the sentinel itself is not interpreted as a public Stream-K
control. The launched baseline symbol does end in `_stream_k`.
## Canonical Operation
- Logical GEMM: `M=37,810, N=5,376, K=14,336`.
- Descriptor GEMM after accepted producer padding: `M=37,824`.
- Inputs: packed NVFP4 E2M1 activation and weight with vector-16 UE4M3 scales.
- Accumulation/scalars: FP32; output: BF16.
- Accepted activation producer: `vortex_native_quantize_swiglu_nvfp4`.
- Image: `sha256:1d340e14cb6fc45ccfdbe63dde8db2a2b3aea94b493702c8a08e1f8d5b4f7b83`.
## Timing
The 20-round isolated sweep alternated AB/BA order in one loaded model.
| Measurement | Baseline | Split-K-1 | Result |
| --- | ---: | ---: | ---: |
| FC2 p50 | `53.618 ms` | `15.636 ms` | `3.43x` |
| Dense throughput p50 | `108.70 TFLOP/s` | `372.74 TFLOP/s` | `3.43x` |
| Producer + FC2 p50 | `76.737 ms` | `38.823 ms` | `49.41%` faster |
The negative-sentinel candidate reproduced the production-equivalent baseline
within noise. Explicit split-K factors `2`, `4`, `8`, and `16` were rejected by
`cublasLtMatmulAlgoCheck` for the tested configuration space.
Complete 20-round block gates replaced only
`block.mlp.fc2.forward_swiglu`:
| Block | Baseline p50 | Candidate p50 | Improvement |
| ---: | ---: | ---: | ---: |
| 0 | `457.891 ms` | `420.536 ms` | `8.16%` |
| 24 | `459.928 ms` | `420.243 ms` | `8.63%` |
| 49 | `458.054 ms` | `417.364 ms` | `8.88%` |
All block outputs were byte-exact against both the paired baseline and retained
traversal.
| Trajectory | Baseline | Candidate | Improvement | Correctness |
| --- | ---: | ---: | ---: | --- |
| Two-step | `45.962 s` | `42.516 s` | `7.50%` | video/audio BF16 exact |
| Canonical 12-step | `278.201 s` | `255.371 s` | `8.21%` | video/audio BF16 exact |
Trajectory values are single paired runs in baseline-then-candidate
order, not repeated medians. Their approximately `38 ms` per-FC2 savings agree
with the alternating isolated and complete-block gates, but the precise
end-to-end percentages retain run-order uncertainty.
The 12-step saving of `22.829 s` matches approximately 600 FC2 invocations
multiplied by the isolated roughly `38 ms` saving.
## Traffic Attribution
The one-pass distinct-data footprint is:
| Category | Bytes |
| --- | ---: |
| Activation QDATA + scales | `305,012,736` |
| Weight QDATA + scales | `43,352,064` |
| Padded BF16 output | `406,683,648` |
| Total | `755,048,448` |
Both schedules issue the same nominal operand requests:
| Requested category | Bytes | Geometric reuse |
| --- | ---: | ---: |
| Weight QDATA + scales | `12,832,210,944` | `296x` |
| Activation QDATA + scales | `12,832,210,944` | about `42.1x` |
| Total L2 operand reads | `25,664,424,960` | unchanged |
The schedule does not eliminate CTA-level rereading. It changes whether those
rereads remain cache-resident:
| NCU metric | Baseline `_stream_k` | Split-K-1 |
| --- | ---: | ---: |
| Main-kernel duration | `55.057 ms` | `17.031 ms` |
| Combined L2 hit rate | `53.32%` | `91.10%` |
| L2 read-miss bytes | `11.762 GB` | `1.909 GB` |
| L2 miss-byte proxy | `12.169 GB` | `2.316 GB` |
| Sysmem traffic proxy | `12.176 GB` | `2.327 GB` |
| Off-chip amplification | `16.12x` | `3.07x` |
| SM throughput | `25.11%` | `81.79%` |
| Tensor-pipe active share | `25.01%` | `82.42%` |
| Eligible warps/scheduler | `0.073` | `0.220` |
| Issue rate | `0.057` | `0.168` |
GB10 does not expose the usual discrete-GPU DRAM byte counters in these
captures. L2 misses and the reported sysmem fill/write sectors are used as the
off-chip proxy.
The `9.853 GB` reduction in L2 read misses is avoided operand rereading. NCU
aggregates the two TMA input descriptors, so it cannot defensibly assign exact
miss-byte totals separately to weights and activations. The request geometry
does prove that each input family accounts for `12.832 GB` of requested reads.
There is no material global partial-accumulator or output-reduction traffic:
- Candidate workspace is zero.
- Neither kernel issues global atomics or global reduction operations.
- No auxiliary reduction kernel is launched.
- Candidate writes exactly one padded BF16 output; baseline writes only
`73,728` bytes more.
- Baseline-only local stack traffic is about `22.35 MB` at L1 and almost none
reaches sysmem.
The baseline's dominant scheduler stall is sleeping. Source-correlated samples
land in `NANOSLEEP.SYNCS` polling around synchronization phase checks. Launch
resources are otherwise the same: 12,432 CTAs, 384 threads/CTA, 168 registers
per thread, 89,088 bytes allocated shared memory, and 25% theoretical
occupancy. The gain therefore comes from traversal locality and reduced
synchronization waiting, not occupancy or a reduction workspace.
## Evidence
- `benchmarks/gb10-fc2-nvfp4-library-sweep-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-block-gate-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-trajectory-2step-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-trajectory-12step-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-baseline-profile-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-splitk1-profile-20260825.json`
- `benchmarks/gb10-fc2-nvfp4-baseline-20260825.csv`
- `benchmarks/gb10-fc2-nvfp4-splitk1-20260825.csv`
- `benchmarks/gb10-fc2-nvfp4-baseline-20260825.ncu-rep`
- `benchmarks/gb10-fc2-nvfp4-splitk1-20260825.ncu-rep`
Production source and configuration were not changed by this study. The
candidate is validated only for the canonical descriptor shape and current
GB10/CUDA-library combination. Promotion still requires shape-specific
`AlgoCheck` with a safe fallback, integration behind the existing
`Nvfp4Linear` boundary, broader shape tests, and deployment validation.

1
research/fc2_nvfp4_scheduling/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Isolated cuBLASLt NVFP4 FC2 scheduling research."""

9
research/fc2_nvfp4_scheduling/candidate_splitk1.json Normal file
View file

@ -0,0 +1,9 @@
{
"algorithm_id": 70,
"tile_id": 20,
"stages_id": 37,
"split_k": 1,
"reduction_scheme": 0,
"custom_option": 0,
"cta_swizzle": 0
}

388
research/fc2_nvfp4_scheduling/fc2_nvfp4_lt.cpp Normal file
View file

@ -0,0 +1,388 @@
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublasLt.h>
#include <cuda_runtime_api.h>
#include <array>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>
namespace py = pybind11;
#define LT_CHECK(call) \
do { \
const cublasStatus_t status_ = (call); \
if (status_ != CUBLAS_STATUS_SUCCESS) \
throw std::runtime_error(std::string(#call) + " failed: " + \
std::to_string(static_cast<int>(status_))); \
} while (0)
namespace {
thread_local cublasLtHandle_t handle = nullptr;
cublasLtHandle_t get_handle() {
if (!handle) LT_CHECK(cublasLtCreate(&handle));
return handle;
}
void require_cuda_contiguous(const torch::Tensor& tensor, const char* name) {
TORCH_CHECK(tensor.is_cuda(), name, " must be CUDA");
TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
}
void require_same_device(const torch::Tensor& tensor,
const torch::Tensor& reference,
const char* name) {
TORCH_CHECK(tensor.device() == reference.device(), name,
" must be on the same CUDA device as activation_qdata");
}
struct Problem {
cublasLtMatmulDesc_t operation = nullptr;
cublasLtMatrixLayout_t a = nullptr, b = nullptr, c = nullptr, d = nullptr;
Problem(const torch::Tensor& activation_qdata,
const torch::Tensor& activation_block_scale,
const torch::Tensor& weight_qdata,
const torch::Tensor& weight_block_scale) {
require_cuda_contiguous(activation_qdata, "activation_qdata");
require_cuda_contiguous(activation_block_scale, "activation_block_scale");
require_cuda_contiguous(weight_qdata, "weight_qdata");
require_cuda_contiguous(weight_block_scale, "weight_block_scale");
require_same_device(activation_block_scale, activation_qdata,
"activation_block_scale");
require_same_device(weight_qdata, activation_qdata, "weight_qdata");
require_same_device(weight_block_scale, activation_qdata,
"weight_block_scale");
TORCH_CHECK(activation_qdata.scalar_type() == at::kByte &&
weight_qdata.scalar_type() == at::kByte,
"packed NVFP4 operands must use uint8 storage");
TORCH_CHECK(activation_qdata.dim() == 2 && weight_qdata.dim() == 2,
"packed NVFP4 operands must be rank two");
TORCH_CHECK(activation_block_scale.element_size() == 1 &&
weight_block_scale.element_size() == 1,
"NVFP4 block scales must use one-byte E4M3 storage");
TORCH_CHECK(activation_qdata.size(1) == weight_qdata.size(1),
"packed K dimensions differ");
// This is deliberately the same column-major reinterpretation used by
// Comfy Kitchen 0.2.31: weight is Lt A, activation is Lt B, and D is D^T.
const int64_t m = weight_qdata.size(0); // row-major N
const int64_t n = activation_qdata.size(0); // row-major padded M
const int64_t k = activation_qdata.size(1) * 2;
TORCH_CHECK(activation_block_scale.numel() >= n * (k / 16),
"activation_block_scale storage is too small");
TORCH_CHECK(weight_block_scale.numel() >= m * (k / 16),
"weight_block_scale storage is too small");
LT_CHECK(cublasLtMatmulDescCreate(&operation, CUBLAS_COMPUTE_32F,
CUDA_R_32F));
cublasLtMatmulMatrixScale_t scale_mode =
CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scale_mode,
sizeof(scale_mode)));
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &scale_mode,
sizeof(scale_mode)));
const cublasOperation_t transa = CUBLAS_OP_T;
const cublasOperation_t transb = CUBLAS_OP_N;
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)));
const void* a_scale = weight_block_scale.data_ptr();
const void* b_scale = activation_block_scale.data_ptr();
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale,
sizeof(a_scale)));
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale,
sizeof(b_scale)));
const cublasDataType_t scale_type = CUDA_R_32F;
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type,
sizeof(scale_type)));
const cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode,
sizeof(pointer_mode)));
const cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
LT_CHECK(cublasLtMatmulDescSetAttribute(
operation, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue,
sizeof(epilogue)));
LT_CHECK(cublasLtMatrixLayoutCreate(&a, CUDA_R_4F_E2M1, k, m, k));
LT_CHECK(cublasLtMatrixLayoutCreate(&b, CUDA_R_4F_E2M1, k, n, k));
LT_CHECK(cublasLtMatrixLayoutCreate(&c, CUDA_R_16BF, m, n, m));
LT_CHECK(cublasLtMatrixLayoutCreate(&d, CUDA_R_16BF, m, n, m));
}
~Problem() {
if (d) cublasLtMatrixLayoutDestroy(d);
if (c) cublasLtMatrixLayoutDestroy(c);
if (b) cublasLtMatrixLayoutDestroy(b);
if (a) cublasLtMatrixLayoutDestroy(a);
if (operation) cublasLtMatmulDescDestroy(operation);
}
};
template <typename T>
bool config_get(const cublasLtMatmulAlgo_t& algo,
cublasLtMatmulAlgoConfigAttributes_t attr, T* value) {
size_t written = 0;
return cublasLtMatmulAlgoConfigGetAttribute(&algo, attr, value,
sizeof(T), &written) ==
CUBLAS_STATUS_SUCCESS &&
written == sizeof(T);
}
template <typename T>
void put_config(py::dict& result, const char* name,
const cublasLtMatmulAlgo_t& algo,
cublasLtMatmulAlgoConfigAttributes_t attr) {
T value{};
if (config_get(algo, attr, &value)) result[name] = value;
}
template <typename T>
void put_cap_scalar(py::dict& caps, const char* name,
const cublasLtMatmulAlgo_t& algo,
cublasLtMatmulAlgoCapAttributes_t attr) {
T value{};
size_t written = 0;
if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, &value, sizeof(value),
&written) == CUBLAS_STATUS_SUCCESS &&
written == sizeof(value))
caps[name] = value;
}
void put_cap_array(py::dict& caps, const char* name,
const cublasLtMatmulAlgo_t& algo,
cublasLtMatmulAlgoCapAttributes_t attr) {
size_t bytes = 0;
if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, nullptr, 0, &bytes) !=
CUBLAS_STATUS_SUCCESS ||
bytes == 0)
return;
std::vector<uint32_t> values((bytes + sizeof(uint32_t) - 1) /
sizeof(uint32_t));
size_t written = 0;
if (cublasLtMatmulAlgoCapGetAttribute(&algo, attr, values.data(), bytes,
&written) != CUBLAS_STATUS_SUCCESS)
return;
values.resize(written / sizeof(uint32_t));
caps[name] = values;
}
py::dict describe(const cublasLtMatmulAlgo_t& algo,
const cublasLtMatmulHeuristicResult_t& checked,
cublasStatus_t api_status) {
py::dict result;
for (const char* name : {"algorithm_id", "tile_id", "stages_id", "split_k",
"reduction_scheme", "custom_option", "cta_swizzle",
"inner_shape", "cluster_shape"})
result[name] = py::none();
put_config<int>(result, "algorithm_id", algo, CUBLASLT_ALGO_CONFIG_ID);
put_config<uint32_t>(result, "tile_id", algo, CUBLASLT_ALGO_CONFIG_TILE_ID);
put_config<uint32_t>(result, "stages_id", algo,
CUBLASLT_ALGO_CONFIG_STAGES_ID);
// CUDA 13 documents SPLITK_NUM as int32_t. Preserve negative library
// sentinel values instead of wrapping them into fictitious huge factors.
put_config<int32_t>(result, "split_k", algo,
CUBLASLT_ALGO_CONFIG_SPLITK_NUM);
put_config<uint32_t>(result, "reduction_scheme", algo,
CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME);
put_config<uint32_t>(result, "custom_option", algo,
CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION);
put_config<uint32_t>(result, "cta_swizzle", algo,
CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING);
#if CUDA_VERSION >= 12000
put_config<uint32_t>(result, "inner_shape", algo,
CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID);
put_config<uint32_t>(result, "cluster_shape", algo,
CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID);
#endif
result["required_workspace_bytes"] = checked.workspaceSize;
result["waves"] = checked.wavesCount;
result["state"] = static_cast<int>(checked.state);
result["api_status"] = static_cast<int>(api_status);
result["valid"] = api_status == CUBLAS_STATUS_SUCCESS &&
checked.state == CUBLAS_STATUS_SUCCESS;
py::dict caps;
put_cap_scalar<int>(caps, "split_k_support", algo,
CUBLASLT_ALGO_CAP_SPLITK_SUPPORT);
put_cap_scalar<uint32_t>(caps, "reduction_scheme_mask", algo,
CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK);
put_cap_scalar<uint32_t>(caps, "cta_swizzle_support", algo,
CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT);
put_cap_scalar<int>(caps, "custom_option_max", algo,
CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX);
put_cap_scalar<int>(caps, "strided_batch_support", algo,
CUBLASLT_ALGO_CAP_STRIDED_BATCH_SUPPORT);
put_cap_scalar<int>(caps, "out_of_place_result_support", algo,
CUBLASLT_ALGO_CAP_OUT_OF_PLACE_RESULT_SUPPORT);
put_cap_array(caps, "tile_ids", algo, CUBLASLT_ALGO_CAP_TILE_IDS);
put_cap_array(caps, "stages_ids", algo, CUBLASLT_ALGO_CAP_STAGES_IDS);
caps["inner_cluster_shape_capability_note"] =
"This CUDA 13 cublasLt.h exposes config IDs but no public capability "
"attributes that enumerate inner/cluster shape IDs.";
result["capabilities"] = caps;
return result;
}
cublasLtMatmulAlgo_t init_algo(int algorithm_id) {
cublasLtMatmulAlgo_t algo{};
LT_CHECK(cublasLtMatmulAlgoInit(
get_handle(), CUBLAS_COMPUTE_32F, CUDA_R_32F, CUDA_R_4F_E2M1,
CUDA_R_4F_E2M1, CUDA_R_16BF, CUDA_R_16BF, algorithm_id, &algo));
return algo;
}
template <typename T>
void maybe_set(cublasLtMatmulAlgo_t* algo, const py::dict& config,
const char* key, cublasLtMatmulAlgoConfigAttributes_t attr) {
if (!config.contains(key) || config[key].is_none()) return;
const T value = config[key].cast<T>();
LT_CHECK(cublasLtMatmulAlgoConfigSetAttribute(algo, attr, &value,
sizeof(value)));
}
cublasLtMatmulAlgo_t configured_algo(const py::dict& config) {
TORCH_CHECK(config.contains("algorithm_id"), "algorithm_id is required");
auto algo = init_algo(config["algorithm_id"].cast<int>());
maybe_set<uint32_t>(&algo, config, "tile_id", CUBLASLT_ALGO_CONFIG_TILE_ID);
maybe_set<uint32_t>(&algo, config, "stages_id",
CUBLASLT_ALGO_CONFIG_STAGES_ID);
maybe_set<int32_t>(&algo, config, "split_k",
CUBLASLT_ALGO_CONFIG_SPLITK_NUM);
maybe_set<uint32_t>(&algo, config, "reduction_scheme",
CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME);
maybe_set<uint32_t>(&algo, config, "custom_option",
CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION);
maybe_set<uint32_t>(&algo, config, "cta_swizzle",
CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING);
#if CUDA_VERSION >= 12000
maybe_set<uint32_t>(&algo, config, "inner_shape",
CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID);
maybe_set<uint32_t>(&algo, config, "cluster_shape",
CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID);
#endif
return algo;
}
py::list enumerate(torch::Tensor activation_qdata,
torch::Tensor activation_block_scale,
torch::Tensor weight_qdata,
torch::Tensor weight_block_scale,
int64_t max_workspace, int requested_count) {
c10::cuda::CUDAGuard guard(activation_qdata.device());
Problem problem(activation_qdata, activation_block_scale, weight_qdata,
weight_block_scale);
cublasLtMatmulPreference_t preference = nullptr;
LT_CHECK(cublasLtMatmulPreferenceCreate(&preference));
LT_CHECK(cublasLtMatmulPreferenceSetAttribute(
preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace,
sizeof(max_workspace)));
std::vector<cublasLtMatmulHeuristicResult_t> found(requested_count);
int returned = 0;
const auto status = cublasLtMatmulAlgoGetHeuristic(
get_handle(), problem.operation, problem.a, problem.b, problem.c,
problem.d, preference, requested_count, found.data(), &returned);
cublasLtMatmulPreferenceDestroy(preference);
LT_CHECK(status);
py::list output;
for (int i = 0; i < returned; ++i)
output.append(describe(found[i].algo, found[i], CUBLAS_STATUS_SUCCESS));
return output;
}
py::dict check(torch::Tensor activation_qdata,
torch::Tensor activation_block_scale,
torch::Tensor weight_qdata,
torch::Tensor weight_block_scale, py::dict config) {
c10::cuda::CUDAGuard guard(activation_qdata.device());
Problem problem(activation_qdata, activation_block_scale, weight_qdata,
weight_block_scale);
auto algo = configured_algo(config);
cublasLtMatmulHeuristicResult_t result{};
const auto status = cublasLtMatmulAlgoCheck(
get_handle(), problem.operation, problem.a, problem.b, problem.c,
problem.d, &algo, &result);
return describe(algo, result, status);
}
void run(torch::Tensor activation_qdata,
torch::Tensor activation_block_scale,
torch::Tensor weight_qdata, torch::Tensor weight_block_scale,
torch::Tensor alpha, torch::Tensor beta, torch::Tensor output,
torch::Tensor workspace, py::dict config) {
c10::cuda::CUDAGuard guard(activation_qdata.device());
require_cuda_contiguous(alpha, "alpha");
require_cuda_contiguous(beta, "beta");
require_cuda_contiguous(output, "output");
require_cuda_contiguous(workspace, "workspace");
require_same_device(alpha, activation_qdata, "alpha");
require_same_device(beta, activation_qdata, "beta");
require_same_device(output, activation_qdata, "output");
require_same_device(workspace, activation_qdata, "workspace");
TORCH_CHECK(alpha.scalar_type() == at::kFloat && alpha.numel() == 1,
"alpha must be one device FP32 value");
TORCH_CHECK(beta.scalar_type() == at::kFloat && beta.numel() == 1,
"beta must be one device FP32 value");
TORCH_CHECK(output.scalar_type() == at::kBFloat16 && output.dim() == 2,
"output must be rank-two BF16");
TORCH_CHECK(workspace.scalar_type() == at::kByte,
"workspace must use uint8 storage");
TORCH_CHECK(output.size(0) == activation_qdata.size(0) &&
output.size(1) == weight_qdata.size(0),
"output must be [packed activation rows, weight rows]");
Problem problem(activation_qdata, activation_block_scale, weight_qdata,
weight_block_scale);
auto algo = configured_algo(config);
cublasLtMatmulHeuristicResult_t checked{};
LT_CHECK(cublasLtMatmulAlgoCheck(get_handle(), problem.operation, problem.a,
problem.b, problem.c, problem.d, &algo,
&checked));
TORCH_CHECK(checked.state == CUBLAS_STATUS_SUCCESS,
"selected algorithm failed AlgoCheck with state ",
static_cast<int>(checked.state));
TORCH_CHECK(checked.workspaceSize <= static_cast<size_t>(workspace.numel()),
"selected algorithm requires ", checked.workspaceSize,
" workspace bytes but caller supplied ", workspace.numel());
const auto stream = at::cuda::getCurrentCUDAStream(
activation_qdata.get_device()).stream();
void* workspace_ptr = workspace.numel() ? workspace.data_ptr() : nullptr;
LT_CHECK(cublasLtMatmul(
get_handle(), problem.operation, alpha.data_ptr(), weight_qdata.data_ptr(),
problem.a, activation_qdata.data_ptr(), problem.b, beta.data_ptr(),
output.data_ptr(), problem.c, output.data_ptr(), problem.d, &algo,
workspace_ptr, workspace.numel(), stream));
}
py::dict build_info() {
py::dict result;
result["cuda_version"] = CUDA_VERSION;
result["cublas_version"] = CUBLAS_VERSION;
result["stream_k_public_control"] = false;
result["stream_k_note"] =
"CUDA 13 cuBLASLt exposes no documented MatmulAlgoConfig attribute "
"that directly selects Stream-K. Negative SPLITK_NUM values returned "
"by heuristics are preserved as undocumented library sentinels, not "
"claimed as public Stream-K control.";
return result;
}
} // namespace
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
module.def("enumerate", &enumerate);
module.def("check", &check);
module.def("run", &run);
module.def("build_info", &build_info);
}

729
tools/benchmark_fc2_nvfp4_algorithms.py Normal file
View file

@ -0,0 +1,729 @@
"""Characterize isolated cuBLASLt NVFP4 scheduling at the real H3 FC2 boundary."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import platform
import statistics
import subprocess
import sys
import time
import types
from pathlib import Path
from typing import Any, Callable
import torch
from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.nvfp4_quant import vortex_native_quantize_swiglu_nvfp4
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.rope import h3_rope_rotation
from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas, sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents
ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / "research" / "fc2_nvfp4_scheduling" / "fc2_nvfp4_lt.cpp"
DEFAULT_BUDGETS = (0, 4 << 20, 8 << 20, 16 << 20, 32 << 20, 64 << 20)
DEFAULT_BLOCKS = (0, 24, 49)
ENV_PREFIXES = ("H3_", "COMFY_KITCHEN_", "CUDA_", "TORCH_")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--mode", choices=("compile", "characterize", "sweep", "selected", "profile", "block-gate", "trajectory"), default="characterize")
parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
parser.add_argument("--capture", type=Path, help="Directory containing input.pt, or an input .pt file")
parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/fc2-nvfp4-scheduling.json"))
parser.add_argument("--build-directory", type=Path)
parser.add_argument("--verbose-build", action="store_true")
parser.add_argument("--width", type=int, default=1344)
parser.add_argument("--height", type=int, default=768)
parser.add_argument("--frames", type=int, default=124)
parser.add_argument("--steps", type=int, default=12)
parser.add_argument("--sampler-step", type=int, default=1)
parser.add_argument("--seed", type=int, default=440420)
parser.add_argument("--text-tokens", type=int, default=100)
parser.add_argument("--blocks", type=int, nargs="+", default=list(DEFAULT_BLOCKS))
parser.add_argument("--probe-block", type=int, default=24)
parser.add_argument("--workspace-budgets", type=int, nargs="+", default=list(DEFAULT_BUDGETS))
parser.add_argument("--requested-count", type=int, default=32)
parser.add_argument("--explicit-split-k", type=int, nargs="+", default=[1, 2, 4, 8, 16])
parser.add_argument("--max-explicit-checks", type=int, default=64)
parser.add_argument("--candidate-limit", type=int, default=8)
parser.add_argument("--candidate", help="JSON object, JSON file, or enumerate:<budget-index>:<result-index>")
parser.add_argument("--profile-target", choices=("baseline", "candidate"), default="candidate")
parser.add_argument("--workspace-bytes", type=int, default=32 << 20)
parser.add_argument("--warmup", type=int, default=2)
parser.add_argument("--rounds", type=int, default=6)
parser.add_argument("--checkpoint-sha256", action="store_true", help="Expensive: actually calculate and record the checkpoint SHA-256")
return parser.parse_args()
def load_extension(args: argparse.Namespace):
from torch.utils.cpp_extension import CUDA_HOME, load
if CUDA_HOME is None:
raise RuntimeError("CUDA_HOME is unavailable; CUDA 12.9 or newer headers are required")
kwargs: dict[str, Any] = {}
if args.build_directory is not None:
args.build_directory.mkdir(parents=True, exist_ok=True)
kwargs["build_directory"] = str(args.build_directory)
return load(
name="h3_fc2_nvfp4_lt_schedule",
sources=[str(SOURCE)],
extra_include_paths=[str(Path(CUDA_HOME) / "include")],
extra_cflags=["-O2", "-std=c++17"],
extra_ldflags=["-L" + str(Path(CUDA_HOME) / "lib64"), "-lcublasLt", "-lcublas", "-lcudart"],
verbose=args.verbose_build,
**kwargs,
)
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(8 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def tensor_sha256(tensor: torch.Tensor) -> str:
immutable = tensor.detach().contiguous().clone().view(torch.uint16).cpu()
return hashlib.sha256(immutable.numpy().tobytes()).hexdigest()
def compare(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, Any]:
actual_copy = actual.detach().clone()
expected_copy = expected.detach().clone()
different = int(torch.count_nonzero(actual_copy != expected_copy))
delta = (actual_copy.float() - expected_copy.float()).abs()
return {
"bf16_exact": different == 0,
"different_elements": different,
"max_abs": float(delta.max()),
"mean_abs": float(delta.mean()),
"actual_sha256": tensor_sha256(actual_copy),
"expected_sha256": tensor_sha256(expected_copy),
}
def percentile(values: list[float], fraction: float) -> float:
ordered = sorted(values)
if len(ordered) == 1:
return ordered[0]
position = (len(ordered) - 1) * fraction
lower = int(position)
weight = position - lower
return ordered[lower] * (1.0 - weight) + ordered[min(lower + 1, len(ordered) - 1)] * weight
def timing_summary(milliseconds: list[float], m: int | None = None, n: int | None = None, k: int | None = None) -> dict[str, Any]:
result: dict[str, Any] = {
"samples_ms": milliseconds,
"p50_ms": percentile(milliseconds, 0.50),
"p95_ms": percentile(milliseconds, 0.95),
"mean_ms": statistics.fmean(milliseconds),
}
if m is not None and n is not None and k is not None:
result["dense_tflop_s_p50"] = 2.0 * m * n * k / (result["p50_ms"] * 1.0e9)
return result
def timed_cuda(call: Callable[[], torch.Tensor]) -> tuple[float, torch.Tensor]:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
output = call()
end.record()
end.synchronize()
return float(start.elapsed_time(end)), output
def environment(args: argparse.Namespace, extension) -> dict[str, Any]:
checkpoint = Path(args.model_path)
try:
commit = subprocess.run(
["git", "rev-parse", "HEAD"], cwd=ROOT, check=True,
capture_output=True, text=True,
).stdout.strip()
except (OSError, subprocess.CalledProcessError):
commit = None
result = {
"platform": platform.platform(),
"python": sys.version,
"torch": torch.__version__,
"torch_cuda": torch.version.cuda,
"device": torch.cuda.get_device_name(),
"device_capability": list(torch.cuda.get_device_capability()),
"driver": torch.cuda.driver_version() if hasattr(torch.cuda, "driver_version") else None,
"git_commit": commit,
"checkpoint_path": str(checkpoint),
"checkpoint_sha256": sha256_file(checkpoint) if args.checkpoint_sha256 else None,
"checkpoint_hash_note": "calculated" if args.checkpoint_sha256 else "not calculated",
"environment_switches": {key: value for key, value in sorted(os.environ.items()) if key.startswith(ENV_PREFIXES)},
"extension": dict(extension.build_info()),
}
try:
import comfy_kitchen
result["comfy_kitchen"] = getattr(comfy_kitchen, "__version__", "0.2.31 package without __version__")
except Exception as error:
result["comfy_kitchen"] = f"import error: {error}"
return result
def make_workload(args: argparse.Namespace, checkpoint: H3Checkpoint, model: H3PackedDenoiser):
torch.manual_seed(args.seed)
if args.capture is not None:
path = args.capture / "input.pt" if args.capture.is_dir() else args.capture
payload = torch.load(path, map_location="cuda", weights_only=False)
hidden = payload["hidden"].to("cuda").contiguous()
timesteps = payload["timesteps"].to("cuda")
positions = payload["position_ids"].to("cuda")
segments = payload["segments"]
metadata = {"capture": str(path)}
else:
packer = H3PromptPacker(checkpoint)
video, audio, aligned_frames = random_av_latents(
args.width, args.height, args.frames, args.seed, device="cuda",
)
sigma = beta_sigmas(args.steps, device="cuda")[args.sampler_step - 1]
native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma)
text = torch.randn(1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16)
hidden, timesteps, segments, positions, _, _ = packer(
text, video, native_audio, _model_sigma(sigma),
)
metadata = {
"resolution": [args.width, args.height],
"frames": aligned_frames,
"steps": args.steps,
"sampler_step": args.sampler_step,
"seed": args.seed,
"text_tokens": args.text_tokens,
}
canonical = (args.width, args.height, args.frames, args.seed, args.sampler_step, args.text_tokens) == (1344, 768, 124, 440420, 1, 100)
if canonical and hidden.shape[0] != 37_810:
raise RuntimeError(f"canonical workload must contain 37,810 tokens, got {hidden.shape[0]}")
rotation = h3_rope_rotation(positions.to("cuda"), model.backbone.inv_freq, torch.bfloat16)
metadata.update({"tokens": hidden.shape[0], "hidden_shape": list(hidden.shape), "segments": segments})
return hidden, timesteps, rotation, segments, metadata
def capture_boundaries(args: argparse.Namespace, model: H3PackedDenoiser, hidden, timesteps, rotation, segments):
wanted = set(args.blocks)
if wanted != set(DEFAULT_BLOCKS):
missing = set(DEFAULT_BLOCKS) - wanted
if missing:
raise ValueError(f"--blocks must retain required blocks 0,24,49; missing {sorted(missing)}")
block_inputs: dict[int, torch.Tensor] = {}
block_outputs: dict[int, torch.Tensor] = {}
gate_up: dict[int, torch.Tensor] = {}
hooks = []
modulated_forwards = {}
for index in wanted:
fc1 = model.backbone.blocks[index].mlp.fc1
hooks.append(fc1.register_forward_hook(
lambda _module, _inputs, output, index=index: gate_up.__setitem__(index, output.detach().clone())
))
original = fc1.forward_modulated
modulated_forwards[index] = original
def capture_modulated(_self, *values, index=index, original=original, **kwargs):
output = original(*values, **kwargs)
gate_up[index] = output.detach().clone()
return output
fc1.forward_modulated = types.MethodType(capture_modulated, fc1)
try:
with torch.inference_mode():
for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)):
if index in wanted:
block_inputs[index] = hidden.detach().clone()
hidden = block(hidden, rotation, *adaln(timesteps), segments)
if index in wanted:
block_outputs[index] = hidden.detach().clone()
finally:
for hook in hooks:
hook.remove()
for index, original in modulated_forwards.items():
model.backbone.blocks[index].mlp.fc1.forward_modulated = original
if set(gate_up) != wanted:
raise RuntimeError(f"failed to capture FC1 boundaries: got {sorted(gate_up)}")
return block_inputs, block_outputs, gate_up
def packed_boundary(gate_up: torch.Tensor, fc2):
tensor_scale, qdata, block_scale = vortex_native_quantize_swiglu_nvfp4(gate_up)
alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
beta = torch.zeros(1, device=gate_up.device, dtype=torch.float32)
return tensor_scale, qdata, block_scale, alpha, beta
def baseline_metadata(fc2, gate_up: torch.Tensor) -> dict[str, Any]:
from torch.profiler import ProfilerActivity, profile
with torch.inference_mode():
fc2.forward_swiglu(gate_up)
torch.cuda.synchronize()
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as captured:
output = fc2.forward_swiglu(gate_up)
torch.cuda.synchronize()
events = []
for event in captured.events():
if "CUDA" not in str(getattr(event, "device_type", "")):
continue
device_us = max(
float(getattr(event, "device_time_total", 0.0) or 0.0),
float(getattr(event, "self_device_time_total", 0.0) or 0.0),
float(getattr(event, "cuda_time_total", 0.0) or 0.0),
float(getattr(event, "self_cuda_time_total", 0.0) or 0.0),
)
events.append({"name": event.name, "device_time_total_us": device_us})
events.sort(key=lambda row: row["device_time_total_us"], reverse=True)
return {
"path": "fc2.forward_swiglu -> accepted producer -> Comfy Kitchen 0.2.31 scaled_mm_nvfp4",
"descriptors": {
"packed_input_output": "row-major [M,K] @ [N,K].T -> BF16 [M,N]",
"block_scale_mode": "VEC16_UE4M3",
"compute_and_scale": "FP32",
"scalar_pointer_mode": "device",
"bias": None,
"beta": 0.0,
"comfy_kitchen_version": "0.2.31",
},
"profiler_cuda_events_available": bool(events),
"profiler_note": None if events else "Torch profiler returned no CUDA kernel events on this build; use --mode profile with NCU for kernel metadata.",
"top_cuda_events": events[:10],
"output_sha256": tensor_sha256(output.detach().clone()),
}
def enumerate_all(extension, packed, fc2, args: argparse.Namespace):
_tensor_scale, qdata, block_scale, _alpha, _beta = packed
groups = []
errors = []
for budget in args.workspace_budgets:
try:
values = [dict(value) for value in extension.enumerate(
qdata, block_scale, fc2.weight, fc2.weight_scale, budget, args.requested_count,
)]
groups.append({"max_workspace_bytes": budget, "requested_count": args.requested_count, "returned_count": len(values), "algorithms": values})
except Exception as error:
groups.append({"max_workspace_bytes": budget, "requested_count": args.requested_count, "returned_count": 0, "algorithms": []})
errors.append({"operation": "heuristic", "max_workspace_bytes": budget, "error": repr(error)})
return groups, errors
def explicit_checks(extension, packed, fc2, enumerated, args: argparse.Namespace):
_tensor_scale, qdata, block_scale, _alpha, _beta = packed
unique_bases: dict[int, dict[str, Any]] = {}
for group in enumerated:
for candidate in group["algorithms"]:
unique_bases.setdefault(candidate["algorithm_id"], candidate)
checked, errors = [], []
attempts = 0
for base in unique_bases.values():
mask = int(base.get("capabilities", {}).get("reduction_scheme_mask", 0))
reductions = [1 << bit for bit in range(32) if mask & (1 << bit)]
for split_k in args.explicit_split_k:
schemes = [0] if split_k == 1 else reductions
if not schemes:
errors.append({"algorithm_id": base["algorithm_id"], "split_k": split_k, "error": "capability reports no public reduction scheme"})
for reduction in schemes:
if attempts >= args.max_explicit_checks:
return checked, errors
attempts += 1
config = {key: base[key] for key in ("algorithm_id", "tile_id", "stages_id", "custom_option", "cta_swizzle", "inner_shape", "cluster_shape") if key in base}
config.update({"split_k": split_k, "reduction_scheme": reduction})
try:
result = dict(extension.check(qdata, block_scale, fc2.weight, fc2.weight_scale, config))
result["requested_config"] = config
checked.append(result)
except Exception as error:
errors.append({"config": config, "error": repr(error)})
return checked, errors
def resolve_candidate(args: argparse.Namespace, enumerated, explicit) -> dict[str, Any]:
if args.candidate:
if args.candidate.startswith("enumerate:"):
_, group, result = args.candidate.split(":")
return dict(enumerated[int(group)]["algorithms"][int(result)])
path = Path(args.candidate)
return json.loads(path.read_text()) if path.exists() else json.loads(args.candidate)
valid_explicit = [row for row in explicit if row.get("valid") and row.get("required_workspace_bytes", sys.maxsize) <= args.workspace_bytes]
if valid_explicit:
return valid_explicit[0]["requested_config"]
for group in enumerated:
for row in group["algorithms"]:
if row.get("valid") and row.get("required_workspace_bytes", sys.maxsize) <= args.workspace_bytes:
return row
raise RuntimeError("no valid candidate fits --workspace-bytes")
def candidate_runner(extension, packed, fc2, config, workspace_bytes: int):
_tensor_scale, qdata, block_scale, alpha, beta = packed
output = torch.empty((qdata.shape[0], fc2.out_features), device=qdata.device, dtype=torch.bfloat16)
workspace = torch.empty(workspace_bytes, device=qdata.device, dtype=torch.uint8)
def run() -> torch.Tensor:
extension.run(qdata, block_scale, fc2.weight, fc2.weight_scale, alpha, beta, output, workspace, config)
return output
return run, output, workspace
def comfy_packed_runner(packed, fc2, logical_rows: int):
import torch.nn.functional as functional
from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout
tensor_scale, qdata, block_scale, _alpha, _beta = packed
activation = QuantizedTensor(
qdata,
"TensorCoreNVFP4Layout",
TensorCoreNVFP4Layout.Params(
scale=tensor_scale,
orig_dtype=torch.bfloat16,
orig_shape=(logical_rows, fc2.in_features),
block_scale=block_scale,
),
)
weight = fc2._packed_weight()
def run():
return functional.linear(activation, weight, None)[:logical_rows, :fc2.out_features]
return run
def benchmark_pair(baseline, candidate, rounds: int, warmup: int, m: int, n: int, k: int, *, report_dense: bool = True):
with torch.inference_mode():
for _ in range(warmup):
baseline()
candidate()
torch.cuda.synchronize()
samples = {"baseline": [], "candidate": []}
last = {}
for round_index in range(rounds):
order = (("baseline", baseline), ("candidate", candidate)) if round_index % 2 == 0 else (("candidate", candidate), ("baseline", baseline))
for name, call in order:
elapsed, output = timed_cuda(call)
samples[name].append(elapsed)
last[name] = output.detach().clone()
return {
"order": "AB/BA alternates by round",
"baseline": timing_summary(samples["baseline"], m if report_dense else None, n if report_dense else None, k if report_dense else None),
"candidate": timing_summary(samples["candidate"], m if report_dense else None, n if report_dense else None, k if report_dense else None),
"parity": compare(last["candidate"][:m, :n], last["baseline"][:m, :n]),
}
def benchmark_fc2_candidate(extension, gate_up, fc2, packed, config, args):
run_candidate, output, workspace = candidate_runner(extension, packed, fc2, config, args.workspace_bytes)
m, n, k = gate_up.shape[0], fc2.out_features, fc2.in_features
run_comfy_packed = comfy_packed_runner(packed, fc2, m)
direct = benchmark_pair(
run_comfy_packed, run_candidate,
args.rounds, args.warmup, m, n, k,
)
boundary_output = torch.empty_like(output)
boundary_workspace = torch.empty_like(workspace)
def candidate_boundary():
fresh = packed_boundary(gate_up, fc2)
extension.run(
fresh[1], fresh[2], fc2.weight, fc2.weight_scale, fresh[3], fresh[4],
boundary_output, boundary_workspace, config,
)
return boundary_output
boundary = benchmark_pair(
lambda: fc2.forward_swiglu(gate_up), candidate_boundary,
args.rounds, args.warmup, m, n, k,
)
if not direct["parity"]["bf16_exact"] or not boundary["parity"]["bf16_exact"]:
raise RuntimeError("FC2 library candidate is not byte-exact")
checked = dict(extension.check(packed[1], packed[2], fc2.weight, fc2.weight_scale, config))
return {
"selected_config": config,
"supplied_workspace_bytes": workspace.numel(),
"required_workspace_bytes": checked.get("required_workspace_bytes"),
"checked": checked,
"fc2_only": direct,
"accepted_producer_plus_fc2": boundary,
"output_sha256": tensor_sha256(output[:m].detach().clone()),
}
def profile_one(extension, packed, fc2, config, logical_rows: int, args):
run_candidate, output, workspace = candidate_runner(extension, packed, fc2, config, args.workspace_bytes)
expected = None
if args.profile_target == "baseline":
run_profiled = comfy_packed_runner(packed, fc2, logical_rows)
else:
run_profiled = run_candidate
with torch.inference_mode():
for _ in range(args.warmup):
run_profiled()
if args.profile_target == "candidate":
expected = comfy_packed_runner(packed, fc2, logical_rows)().detach().clone()
torch.cuda.synchronize()
torch.cuda.cudart().cudaProfilerStart()
output = run_profiled()
torch.cuda.cudart().cudaProfilerStop()
torch.cuda.synchronize()
parity = compare(output[:logical_rows], expected[:logical_rows]) if expected is not None else None
if parity is not None and not parity["bf16_exact"]:
raise RuntimeError("profiled FC2 library candidate is not byte-exact")
return {
"target": args.profile_target,
"selected_config": config if args.profile_target == "candidate" else None,
"supplied_workspace_bytes": workspace.numel() if args.profile_target == "candidate" else 32 << 20,
"output_sha256": tensor_sha256(output[:logical_rows, :fc2.out_features]),
"candidate_vs_baseline": parity,
}
def block_gate(extension, model, block_inputs, block_outputs, gate_up, timesteps, rotation, segments, config, args):
rows = []
for index in DEFAULT_BLOCKS:
block = model.backbone.blocks[index]
fc2 = block.mlp.fc2
packed = packed_boundary(gate_up[index], fc2)
_tensor_scale, qdata, _block_scale, _alpha, beta = packed
candidate_output = torch.empty(
(qdata.shape[0], fc2.out_features), device=qdata.device, dtype=torch.bfloat16,
)
workspace = torch.empty(args.workspace_bytes, device=qdata.device, dtype=torch.uint8)
original = fc2.forward_swiglu
def replacement(_self, actual_gate_up, expected=gate_up[index]):
if actual_gate_up.shape != expected.shape:
raise RuntimeError("block-gate FC2 received an unexpected boundary shape")
tensor_scale, actual_qdata, actual_block_scale = vortex_native_quantize_swiglu_nvfp4(actual_gate_up)
alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
extension.run(
actual_qdata, actual_block_scale, fc2.weight, fc2.weight_scale,
alpha, beta, candidate_output, workspace, config,
)
return candidate_output[:actual_gate_up.shape[0], :fc2.out_features]
candidate_method = types.MethodType(replacement, fc2)
adaln_values = tuple(value.detach().clone() for value in model.backbone.adaln[index](timesteps))
def baseline():
fc2.forward_swiglu = original
return block(block_inputs[index].detach().clone(), rotation, *adaln_values, segments)
def candidate():
fc2.forward_swiglu = candidate_method
return block(block_inputs[index].detach().clone(), rotation, *adaln_values, segments)
try:
timing = benchmark_pair(
baseline, candidate, args.rounds, args.warmup,
block_outputs[index].shape[0], block_outputs[index].shape[1], 1,
report_dense=False,
)
baseline_value = baseline().detach().clone()
candidate_value = candidate().detach().clone()
finally:
fc2.forward_swiglu = original
row = {
"block": index,
"only_monkeypatched_method": "block.mlp.fc2.forward_swiglu",
"accepted_gate_and_residual_path_preserved": True,
"candidate_vs_baseline": compare(candidate_value, baseline_value),
"baseline_vs_traversal": compare(baseline_value, block_outputs[index]),
"candidate_vs_traversal": compare(candidate_value, block_outputs[index]),
"timing": timing,
"supplied_workspace_bytes": workspace.numel(),
}
if not all(
row[name]["bf16_exact"]
for name in ("candidate_vs_baseline", "baseline_vs_traversal", "candidate_vs_traversal")
):
raise RuntimeError(f"FC2 library candidate is not byte-exact in block {index}")
rows.append(row)
return rows
def trajectory_gate(extension, checkpoint, model, config, args):
packer = H3PromptPacker(checkpoint)
torch.manual_seed(args.seed)
video, audio, aligned_frames = random_av_latents(
args.width, args.height, args.frames, args.seed, device="cuda",
)
text = torch.randn(
1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16,
)
originals = [block.mlp.fc2.forward_swiglu for block in model.backbone.blocks]
shared_output = torch.empty(
(((37_810 + 15) // 16) * 16, 5376), device="cuda", dtype=torch.bfloat16,
)
workspace = torch.empty(args.workspace_bytes, device="cuda", dtype=torch.uint8)
beta = torch.zeros(1, device="cuda", dtype=torch.float32)
def candidate_method(fc2):
def replacement(_self, actual_gate_up):
tensor_scale, qdata, block_scale = vortex_native_quantize_swiglu_nvfp4(actual_gate_up)
if qdata.shape[0] > shared_output.shape[0] or fc2.out_features > shared_output.shape[1]:
raise RuntimeError("trajectory FC2 boundary exceeds the preallocated canonical output")
alpha = (tensor_scale.float() * fc2.weight_scale_2.float()).reshape(1).contiguous()
output = shared_output[:qdata.shape[0], :fc2.out_features]
extension.run(
qdata, block_scale, fc2.weight, fc2.weight_scale,
alpha, beta, output, workspace, config,
)
return output[:actual_gate_up.shape[0], :fc2.out_features]
return types.MethodType(replacement, fc2)
candidates = [candidate_method(block.mlp.fc2) for block in model.backbone.blocks]
def run(candidate: bool):
for index, block in enumerate(model.backbone.blocks):
block.mlp.fc2.forward_swiglu = candidates[index] if candidate else originals[index]
torch.cuda.synchronize()
started = time.perf_counter()
result = sample_video_res_multistep(
model,
packer,
text,
video.detach().clone(),
audio.detach().clone(),
steps=args.steps,
seed=args.seed,
return_audio=True,
progress=True,
)
torch.cuda.synchronize()
return result, time.perf_counter() - started
try:
with torch.inference_mode():
(reference_video, reference_audio), baseline_seconds = run(False)
(candidate_video, candidate_audio), candidate_seconds = run(True)
finally:
for block, original in zip(model.backbone.blocks, originals, strict=True):
block.mlp.fc2.forward_swiglu = original
video_parity = compare(candidate_video, reference_video)
audio_parity = compare(candidate_audio, reference_audio)
result = {
"steps": args.steps,
"resolution": [args.width, args.height],
"frames": aligned_frames,
"seed": args.seed,
"baseline_seconds": baseline_seconds,
"candidate_seconds": candidate_seconds,
"improvement_percent": (1.0 - candidate_seconds / baseline_seconds) * 100.0,
"video_parity": video_parity,
"audio_parity": audio_parity,
"bf16_exact": video_parity["bf16_exact"] and audio_parity["bf16_exact"],
"selected_config": config,
"supplied_workspace_bytes": workspace.numel(),
"all_50_fc2_calls_replaced": True,
"accepted_swiglu_producer_preserved": True,
"accepted_gate_and_residual_path_preserved": True,
}
if not result["bf16_exact"]:
raise RuntimeError("FC2 library candidate trajectory is not byte-exact")
return result
def main() -> None:
args = parse_args()
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required")
extension = load_extension(args)
if args.mode == "compile":
result = {"mode": "compile", "extension": dict(extension.build_info())}
print(json.dumps(result, indent=2), flush=True)
return
checkpoint = H3Checkpoint(args.model_path, device="cuda")
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
hidden, timesteps, rotation, segments, workload = make_workload(args, checkpoint, model)
with torch.inference_mode():
block_inputs, block_outputs, gate_ups = capture_boundaries(
args, model, hidden, timesteps, rotation, segments,
)
probe = args.probe_block
if probe not in gate_ups:
raise ValueError("--probe-block must be one of the retained blocks")
fc2 = model.backbone.blocks[probe].mlp.fc2
if fc2.bias is not None:
raise RuntimeError("this no-bias FC2 scheduling study refuses a biased module")
packed = packed_boundary(gate_ups[probe], fc2)
enumerated, errors = enumerate_all(extension, packed, fc2, args)
explicit, explicit_errors = explicit_checks(extension, packed, fc2, enumerated, args)
errors.extend(explicit_errors)
result: dict[str, Any] = {
"mode": args.mode,
"environment": environment(args, extension),
"workload": workload,
"retained_blocks": list(DEFAULT_BLOCKS),
"immutable_cloned_block_inputs": {str(index): list(value.shape) for index, value in block_inputs.items()},
"fc2_boundary": {
"block": probe,
"gate_up_shape": list(gate_ups[probe].shape),
"activation_qdata_shape": list(packed[1].shape),
"weight_qdata_shape": list(fc2.weight.shape),
"logical_mnk": [gate_ups[probe].shape[0], fc2.out_features, fc2.in_features],
"descriptor_mnk_after_padding": [packed[1].shape[0], fc2.out_features, fc2.in_features],
"producer": "vortex_native_quantize_swiglu_nvfp4",
"no_bias": fc2.bias is None,
},
"baseline_kernel_metadata": baseline_metadata(fc2, gate_ups[probe]),
"heuristics": enumerated,
"explicit_split_k_checks": explicit,
}
if args.mode in {"characterize", "sweep", "selected", "profile", "block-gate", "trajectory"}:
config = resolve_candidate(args, enumerated, explicit)
result["selected"] = config
if args.mode == "characterize":
result["candidate_probe"] = benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, config, args)
elif args.mode == "sweep":
candidates = []
seen = set()
pool = [row for group in enumerated for row in group["algorithms"]] + [row for row in explicit if row.get("valid")]
for row in pool:
candidate = row.get("requested_config", row)
key = tuple(candidate.get(name) for name in ("algorithm_id", "tile_id", "stages_id", "split_k", "reduction_scheme", "custom_option", "cta_swizzle", "inner_shape", "cluster_shape"))
if key in seen or row.get("required_workspace_bytes", 0) > args.workspace_bytes:
continue
seen.add(key)
try:
candidates.append(benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, candidate, args))
except Exception as error:
errors.append({"config": candidate, "operation": "benchmark", "error": repr(error)})
if len(candidates) >= args.candidate_limit:
break
result["candidates"] = candidates
elif args.mode == "selected":
result["candidate_probe"] = benchmark_fc2_candidate(extension, gate_ups[probe], fc2, packed, config, args)
elif args.mode == "profile":
result["profile"] = profile_one(extension, packed, fc2, config, gate_ups[probe].shape[0], args)
elif args.mode == "block-gate":
result["block_gate"] = block_gate(extension, model, block_inputs, block_outputs, gate_ups, timesteps, rotation, segments, config, args)
elif args.mode == "trajectory":
result["trajectory"] = trajectory_gate(extension, checkpoint, model, config, args)
result["errors_and_unsupported"] = errors + [{
"feature": "Stream-K",
"supported_public_control": False,
"reason": extension.build_info()["stream_k_note"],
}]
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result, indent=2), flush=True)
if __name__ == "__main__":
main()