nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[misc]: extend the comment sweep to every file

Browse Source
This commit is contained in:
Aryan Kumar
2026-11-10 02:07:14 +00:00
committed by fv-superagent
parent 74e3061646
commit f496e4b838
155 changed files with 65 additions and 752 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/skills/dreamverse-deploy/scripts/dreamverse-deploy.sh
-19
View File
@@ -77,9 +77,6 @@ if [[ "${1:-}" == "--stop" ]]; then
exit 0
fi
# ---------------------------------------------------------------------------
# Args
# ---------------------------------------------------------------------------
usage() {
cat <<USAGE
@@ -168,9 +165,6 @@ esac
REPO_ROOT="${DREAMVERSE_REPO_ROOT:-$(git rev-parse --show-toplevel 2>/dev/null || pwd)}"
LOG_DIR="${DREAMVERSE_LOG_DIR:-/tmp/opencode/dreamverse-deploy}"
# ---------------------------------------------------------------------------
# Prereq checks
# ---------------------------------------------------------------------------
bail() { echo "error: $*" >&2; exit 3; }
@@ -234,9 +228,6 @@ echo " python: ${CONDA_ENV_PYTHON}"
mkdir -p "${LOG_DIR}"
# ---------------------------------------------------------------------------
# Teardown anything on target ports
# ---------------------------------------------------------------------------
echo "[1/8] killing any existing deploy on ports ${BACKEND_PORT}/${FRONTEND_PORT} and GPU ${GPU}..."
@@ -281,9 +272,6 @@ done
gpu_mem="$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sed -n "$((GPU + 1))p" || echo 0)"
echo " ports cleared; GPU${GPU} at ${gpu_mem} MiB"
# ---------------------------------------------------------------------------
# Launch backend
# ---------------------------------------------------------------------------
echo "[2/8] launching backend on GPU ${GPU} port ${BACKEND_PORT} (warmup=${WARMUP} torch_compile=${TORCH_COMPILE} nvenc=${NVENC})..."
@@ -367,9 +355,6 @@ fi
echo "[4/8] backend /readyz OK"
# ---------------------------------------------------------------------------
# Launch frontend
# ---------------------------------------------------------------------------
echo "[5/8] launching frontend on port ${FRONTEND_PORT}..."
@@ -403,7 +388,6 @@ fi
echo " frontend pid=${frontend_pid} log=${frontend_log}"
# Poll FE root
echo "[6/8] polling http://127.0.0.1:${FRONTEND_PORT}/ ..."
fe_ready=0
for i in $(seq 1 30); do
@@ -428,9 +412,6 @@ fi
echo "[7/8] frontend / OK"
# ---------------------------------------------------------------------------
# Print summary
# ---------------------------------------------------------------------------
cwd="$(readlink "/proc/${backend_pid}/cwd" 2>/dev/null || echo unknown)"
gpu_mem_now="$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sed -n "$((GPU + 1))p" || echo 0)"
.buildkite/scripts/pr_test.sh
-2
View File
@@ -43,7 +43,6 @@ if ! python3 -m modal --version &> /dev/null; then
# --break-system-packages preserves prior `pip install --user` semantics on PEP 668 agents.
uv pip install --system --break-system-packages modal
# Verify installation
if ! python3 -m modal --version &> /dev/null; then
log "Error: Failed to install modal. Please install it manually."
exit 1
@@ -52,7 +51,6 @@ fi
log "modal version: $(python3 -m modal --version)"
# Set up Modal authentication using Buildkite secrets
log "Setting up Modal authentication from Buildkite secrets..."
MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
.github/mergify.yml
-8
View File
@@ -33,9 +33,7 @@ merge_protections:
pull_request_rules:
# ============================================================
# Type labels (from PR title prefix)
# ============================================================
- name: "label type: feat"
conditions:
@@ -117,9 +115,7 @@ pull_request_rules:
label:
add: ["type: skill"]
# ============================================================
# Scope labels (from changed files)
# ============================================================
- name: "label scope: training"
conditions:
@@ -227,9 +223,7 @@ pull_request_rules:
label:
add: ["scope: model"]
# ============================================================
# Merge conflict detection
# ============================================================
- name: label conflicting PRs
conditions:
@@ -249,9 +243,7 @@ pull_request_rules:
label:
remove: [needs-rebase]
# ============================================================
# Auto-merge
# ============================================================
- name: auto-merge when ready and all checks pass
conditions:
.pre-commit-config.yaml
+1 -5
View File
@@ -24,7 +24,7 @@ repos:
require_serial: true
args: [--in-place, --verbose]
language_version: python3.12
additional_dependencies: [toml] # TODO: Remove when yapf is upgraded
additional_dependencies: [toml]
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.11.12
hooks:
@@ -36,10 +36,6 @@ repos:
- id: codespell
additional_dependencies: ['tomli']
args: ['--toml', 'pyproject.toml']
# - repo: https://github.com/PyCQA/isort
# rev: 6.0.1
# hooks:
# - id: isort
- repo: https://github.com/jackdewinter/pymarkdown
rev: v0.9.30
hooks:
apps/dreamverse/dreamverse/gpu_pool.py
+1 -1
View File
@@ -307,7 +307,7 @@ def gpu_worker_process(
message=str(e),
))
return True # Continue loop
return True
if first_cmd is not None:
if first_cmd.type == CommandType.SHUTDOWN:
apps/dreamverse/dreamverse/session/controller.py
-1
View File
@@ -1704,7 +1704,6 @@ class SessionController:
and slot.shared_stream_buffer is not None):
start = chunk_offset
end = start + chunk_length
# Copy out of shared buffer for websocket send.
chunk = bytes(slot.shared_stream_buffer[start:end])
else:
chunk = chunk_bytes or b""
apps/dreamverse/scripts/install_native_ffmpeg.sh
-5
View File
@@ -49,7 +49,6 @@
# toolchain.
set -euo pipefail
# ─── Defaults (override via env) ──────────────────────────────────────────
INSTALL_PREFIX="${INSTALL_PREFIX:-$HOME/opt/ffmpeg-native}"
SOURCE_DIR="${SOURCE_DIR:-$HOME/src/ffmpeg-native}"
X264_REF="${X264_REF:-stable}"
@@ -137,7 +136,6 @@ export CC CXX
[[ -n "${AS:-}" ]] && export AS
echo "[install_native_ffmpeg] toolchain: CC=$CC CXX=$CXX AS=${AS:-<gnu-as>} (uname -m=$ARCH, source: $toolchain_source)"
# ─── Step 0: probe required tools ─────────────────────────────────────────
required=("$CC" "$CXX" make pkg-config git)
[[ "$(uname -m)" == "x86_64" ]] && required+=(nasm)
missing=()
@@ -155,7 +153,6 @@ if (( ${#missing[@]} > 0 )); then
exit 1
fi
# ─── Step 0.5: destructive-path guards ────────────────────────────────────
guard_path() {
local name="$1" value="$2"
case "$value" in
@@ -172,7 +169,6 @@ guard_path() {
guard_path INSTALL_PREFIX "$INSTALL_PREFIX"
guard_path SOURCE_DIR "$SOURCE_DIR"
# ─── Step 1: clean ────────────────────────────────────────────────────────
echo "[install_native_ffmpeg] cleaning prior install + sources"
rm -rf "$INSTALL_PREFIX" "$SOURCE_DIR/x264" "$SOURCE_DIR/ffmpeg" "$SOURCE_DIR/nv-codec-headers"
mkdir -p "$SOURCE_DIR" "$INSTALL_PREFIX/lib"
@@ -271,7 +267,6 @@ git clone --depth 1 --branch "$FFMPEG_REF" \
make install
)
# ─── Step 4: sanity check ──────────────────────────────────────────────────
ffmpeg_bin="$INSTALL_PREFIX/bin/ffmpeg"
echo "[install_native_ffmpeg] verifying $ffmpeg_bin"
"$ffmpeg_bin" -hide_banner -buildconf | grep -i -E 'libx264|lto'
apps/dreamverse/scripts/launch/launch_demo.sh
-2
View File
@@ -96,7 +96,6 @@ wait_for() {
return 1
}
# --- backend ---
echo "[launch-demo] backend logs: ${BE_LOG}"
( "${BE_SCRIPT}" "${BE_PORT_ARGS[@]}" >"${BE_LOG}" 2>&1 ) &
BE_PID=$!
@@ -117,7 +116,6 @@ if [[ "${HEALTH_PATH}" != "${READY_PATH}" ]]; then
fi
fi
# --- frontend ---
if [[ "${NO_FRONTEND:-0}" != "1" ]]; then
echo "[launch-demo] frontend logs: ${FE_LOG}"
( "${SCRIPT_DIR}/launch_frontend.sh" >"${FE_LOG}" 2>&1 ) &
apps/dreamverse/web/src/app/globals.css
-5
View File
@@ -110,7 +110,6 @@
--radius-xl: calc(var(--radius) + 4px);
}
/* ——— Resets ——— */
*,
*::before,
@@ -194,7 +193,6 @@ summary::-webkit-details-marker {
color: #ffffff;
}
/* ——— Base layer ——— */
@layer base {
* {
@@ -205,7 +203,6 @@ summary::-webkit-details-marker {
}
}
/* ——— Utilities ——— */
@layer utilities {
.scrollbar-hidden {
@@ -269,7 +266,6 @@ summary::-webkit-details-marker {
animation: stt-pulse 1.4s ease-in-out infinite;
}
/* ——— Theme transition ——— */
html.theme-transition,
html.theme-transition *,
@@ -284,7 +280,6 @@ html.theme-transition *::after {
stroke 300ms ease !important;
}
/* —————————————— CUSTOM TAILWIND —————————————— */
@layer components {
.debug {
apps/dreamverse/web/src/app/page.tsx
-1
View File
@@ -2853,7 +2853,6 @@ export default function Page() {
onSelectCurrent={() => {
const clips = completedClips as any[];
if (sessionStarted && !sessionExpired) {
// Return to live playback
streamStore.patch({ activeClipId: "", activePlaybackStartTime: 0 });
} else if (clips.length > 0) {
selectClip(clips[clips.length - 1].id, 0);
apps/dreamverse/web/src/lib/media/avPipeline.test.ts
-5
View File
@@ -219,15 +219,12 @@ describe('createAvPipeline', () => {
updating = false;
addEventListener(): void {
// noop for test
}
removeEventListener(): void {
// noop for test
}
appendBuffer(): void {
// noop for test
}
}
@@ -335,11 +332,9 @@ describe('createAvPipeline', () => {
appendBuffer = vi.fn();
addEventListener(): void {
// noop for test
}
removeEventListener(): void {
// noop for test
}
}
apps/dreamverse/web/src/lib/media/avPipeline.ts
-1
View File
@@ -270,7 +270,6 @@ export function createAvPipeline({
return;
}
stallRecoveryTimerId = setInterval(() => {
// Stop polling once playback is running normally.
if (initialPlaybackStarted && !playbackResumePending) {
clearStallRecovery();
return;
apps/fastvideo_studio/e2e/datasets.spec.ts
-1
View File
@@ -10,7 +10,6 @@ test.describe('datasets', () => {
skipWithoutMock();
test('lists the seeded datasets', async ({ page, request }) => {
// Read the seeded names from the mock so the assertion tracks the fixture.
const res = await request.get(`${API_BASE}/datasets`);
const datasets = (await res.json()) as Array<{ name: string }>;
expect(datasets.length).toBeGreaterThan(0);
apps/fastvideo_studio/job_runner.py
+2 -2
View File
@@ -711,8 +711,8 @@ class JobRunner:
job.finished_at = None
job._stop_event.clear()
job._log_buf = JobLogBuffer() # fresh log buffer
job.log_file_handler = None # Reset file handler
job._process = None # Reset subprocess handle
job.log_file_handler = None
job._process = None
# Wrap _run_job in an additional safety layer to catch any exceptions
# that might escape (though they shouldn't with our comprehensive handling)
apps/fastvideo_studio/src/app/globals.css
-2
View File
@@ -120,7 +120,6 @@
--radius-xl: calc(var(--radius) + 4px);
}
/* ——— Resets ——— */
html,
body {
@@ -211,7 +210,6 @@ html.theme-transition *::after {
stroke 300ms ease !important;
}
/* ——— Base layer ——— */
@layer base {
button,
apps/fastvideo_studio/src/components/datasets/DatasetSidebar.test.tsx
-3
View File
@@ -54,15 +54,12 @@ describe('DatasetSidebar', () => {
it('lists dataset files after loading', async () => {
render(<DatasetSidebar dataset={dataset} onClose={() => {}} />);
// The first file's caption is rendered once loading resolves.
expect(await screen.findByDisplayValue('cap a')).toBeInTheDocument();
expect(mockedApi.getDatasetFiles).toHaveBeenCalledWith('ds-1');
// One caption editor per returned file.
const captionFields = screen.getAllByPlaceholderText('Caption');
expect(captionFields).toHaveLength(2);
// Media URLs are requested per visible file.
expect(mockedApi.getDatasetMediaUrl).toHaveBeenCalledWith('ds-1', 'a.mp4');
expect(mockedApi.getDatasetMediaUrl).toHaveBeenCalledWith('ds-1', 'b.mp4');
});
apps/fastvideo_studio/src/components/jobs/CreateJobModal.test.tsx
-1
View File
@@ -23,7 +23,6 @@ const MODELS = [
];
beforeEach(() => {
// Reset the shared options store to a known baseline for test isolation.
defaultOptionsStore.set({ options: DEFAULT_OPTIONS });
vi.mocked(getModels).mockResolvedValue(MODELS);
vi.mocked(getDatasets).mockResolvedValue([]);
apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx
-1
View File
@@ -302,7 +302,6 @@ export default function CreateJobModal({
}
}, [isOpen, workloadType, inferenceWorkload, options, editingJobId]);
// Load the models available for this workload.
React.useEffect(() => {
if (!isOpen) return;
// Ignore a superseded response so a slow fetch for a previous workload
apps/fastvideo_studio/src/components/jobs/JobDetailsSidebar.test.tsx
-1
View File
@@ -81,7 +81,6 @@ describe('JobDetailsSidebar', () => {
/>,
);
// Flush the immediate poll fired on mount.
await act(async () => {
await vi.advanceTimersByTimeAsync(0);
});
apps/fastvideo_studio/src/components/jobs/JobQueue.tsx
-2
View File
@@ -101,12 +101,10 @@ export default function JobQueue({ jobType, jobTypesForList }: JobQueueProps) {
fetchJobsRef.current = fetchJobs;
}, [fetchJobs]);
// Initial fetch + refetch whenever an external refresh is triggered.
React.useEffect(() => {
fetchJobs();
}, [fetchJobs, nonce]);
// Poll every second while any job is running/pending; stop otherwise.
const hasActive = jobs.some(
(j) => j.status === 'running' || j.status === 'pending',
);
comfyui/web/js/FastVideo.core.js
+1 -17
View File
@@ -68,7 +68,6 @@ function drawAutoAnnotated(ctx, node, widget_width, y, H) {
ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
ctx.fill();
// Reset fill style for the teeth
if (isAuto) {
ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
} else {
@@ -200,7 +199,6 @@ function mouseAutoAnnotated(event, [x, y], node) {
return true;
}
// Handle increment/decrement buttons if not in auto mode
if (delta !== 0 && !this.isAuto) {
if (this.config[0] === "FVAUTOCOMBO") {
const options = this.options.values || [];
@@ -241,7 +239,6 @@ function mouseAutoAnnotated(event, [x, y], node) {
v += increment;
// Apply min/max constraints
if (this.options.min != null) {
v = Math.max(this.options.min, v);
}
@@ -272,7 +269,6 @@ function mouseAutoAnnotated(event, [x, y], node) {
if (this.config[0] === "FVAUTOCOMBO") {
const options = this.options.values || [];
// Create menu items
const menuItems = options.map(opt => {
const value = typeof opt === 'object' ? opt.value : opt;
const label = typeof opt === 'object' ? opt.label : opt.toString();
@@ -323,7 +319,6 @@ function mouseAutoAnnotated(event, [x, y], node) {
const d_callback = (v) => {
this.value = this.parseValue?.(v) ?? Number(v);
// Apply min/max constraints
if (this.options.min != null) {
this.value = Math.max(this.options.min, this.value);
}
@@ -370,7 +365,6 @@ function makeAutoAnnotated(widget, inputData) {
value: widget.value
};
// Add AUTO properties to the widget
Object.assign(widget, {
type: "BOOLEAN",
draw: drawAutoAnnotated,
@@ -381,7 +375,7 @@ function makeAutoAnnotated(widget, inputData) {
cachedValue: widget.value,
config: inputData,
options: Object.assign({}, inputData[1], widget.options),
original: original, // Store original properties for reference
original: original,
// Disable other potential mouse handlers with no-op functions
onClick: function () {
@@ -429,7 +423,6 @@ function makeAutoAnnotated(widget, inputData) {
return v;
},
serializeValue: function () {
// Return special value for AUTO mode
return this.isAuto ? -99999 : this.value;
},
deserializeValue: function (data) {
@@ -473,28 +466,22 @@ app.registerExtension({
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData?.name == "VideoGenerator" || nodeData?.name === "InferenceArgs" || nodeData?.name === "VAEConfig" ||
nodeData?.name === "TextEncoderConfig" || nodeData?.name === "DITConfig") {
// Add serialization support
chainCallback(nodeType.prototype, "onSerialize", function (info) {
if (!this.widgets) {
return;
}
// Ensure widgets_values exists
if (!info.widgets_values) {
info.widgets_values = {};
}
// Store AUTO widget states in a separate property
if (!info.auto_widget_states) {
info.auto_widget_states = {};
}
// Handle AUTO widgets specially
for (const w of this.widgets) {
if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
// Store the serialized value (for Python node)
info.widgets_values[w.name] = w.serializeValue();
// Store the full state (for UI restoration)
info.auto_widget_states[w.name] = {
isAuto: w.isAuto,
value: w.value,
@@ -504,7 +491,6 @@ app.registerExtension({
}
});
// Add deserialization support
chainCallback(nodeType.prototype, "onConfigure", function (info) {
if (!this.widgets) {
return;
@@ -522,7 +508,6 @@ app.registerExtension({
}
}
// Then, restore full state if available
if (info.auto_widget_states) {
for (const w of this.widgets) {
if (w.type === "BOOLEAN" && w.isAuto !== undefined && w.name in info.auto_widget_states) {
@@ -544,7 +529,6 @@ app.registerExtension({
// Override addInput to handle AUTO widgets
chainCallback(nodeType.prototype, "onNodeCreated", function () {
// Convert any existing widgets to AUTO widgets if needed
let new_widgets = [];
const intWidgetNames = ["sp_size", "tp_size", "height", "width", "num_frames", "num_inference_steps", "flow_shift", "seed", "fps", "scale_factor",
"tile_sample_min_height", "tile_sample_min_width", "tile_sample_min_num_frames", "tile_sample_stride_height", "tile_sample_stride_width",
docs/cookbook/recipes/fastwan21.yaml
-3
View File
@@ -1,6 +1,3 @@
# TODO(cookbook-review): Design sample only. Before merge, double-check the
# selected options, schema types/ranges, and baseline values for each deployment.
# TODO: Recheck field paths and choices when the typed API in PR #1913 is finalized.
title: FastWan2.1 1.3B
summary: Distilled text-to-video serving with sparse attention.
docs/cookbook/recipes/wan21-i2v.yaml
-3
View File
@@ -1,6 +1,3 @@
# TODO(cookbook-review): Design sample only. Before merge, double-check the
# selected options, schema types/ranges, and baseline values for each deployment.
# TODO: Recheck field paths and choices when the typed API in PR #1913 is finalized.
title: Wan2.1 I2V 14B
summary: Image-conditioned video generation from a source frame.
docs/generate_examples.py
+2 -36
View File
@@ -14,7 +14,7 @@ ROOT_DIR = Path(__file__).parent.parent.resolve()
ROOT_DIR_RELATIVE = '../..'
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/getting_started/examples"
GITHUB_REPO = "hao-ai-lab/FastVideo" # Update this to your repo
GITHUB_REPO = "hao-ai-lab/FastVideo"
COOKBOOK_DATA = ROOT_DIR / "docs/assets/cookbook-recipes.json"
COOKBOOK_CLIENTS = {
"python": ("video.py", "python -m pip install openai==3.6.0"),
@@ -364,11 +364,8 @@ class Index:
content = f"# {self.title}\n\n{self.description}\n\n"
if self.caption:
content += f"## {self.caption}\n\n"
# Generate a simple list of documentation links
for doc in self.documents:
# Convert document path to proper link
doc_link = doc.replace("\\", "/")
# Get just the filename for the link text
doc_title = fix_case(Path(doc).stem.replace("_", " ").title())
content += f"- [{doc_title}]({doc_link}.md)\n"
content += "\n"
@@ -459,12 +456,10 @@ class Example:
return fix_case(self.path.stem.replace("_", " ").title())
def generate(self) -> str:
# Create GitHub link to source
github_path = str(self.path.relative_to(ROOT_DIR)).replace("\\", "/")
github_url = f"https://github.com/{GITHUB_REPO}/blob/main/{github_path}"
content = f"**Source:** [{github_path}]({github_url})\n\n"
# Add title for code files
if self.main_file.suffix != ".md":
content += f"# {self.title}\n\n"
@@ -484,19 +479,16 @@ class Example:
return content
content += "## Additional Files\n\n"
# Define binary/non-text file extensions to skip
binary_extensions = {
'.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png', '.webp', '.bmp', '.pdf', '.zip', '.tar',
'.gz', '.mp3', '.wav'
}
for file in sorted(self.other_files):
# Skip binary files
if file.suffix.lower() in binary_extensions:
continue
file_rel_path = file.relative_to(self.path)
# Use collapsible admonitions, normalized by the site preparation
content += f"??? note \"{file_rel_path}\"\n\n"
try:
@@ -517,7 +509,6 @@ class Example:
content += " ```\n"
content += "\n"
except UnicodeDecodeError:
# Skip files that can't be decoded as UTF-8
continue
return content
@@ -580,7 +571,6 @@ def create_category_indices() -> dict[str, Index]:
),
}
# Ensure all category doc directories exist
for index in category_indices.values():
if not index.path.parent.exists():
index.path.parent.mkdir(parents=True)
@@ -604,7 +594,6 @@ def find_examples(category_indices: dict[str, Index], generate_main_index: bool)
dir_name = category_dir_mapping.get(category, category)
category_dir = EXAMPLE_DIR / dir_name
# Skip if directory doesn't exist
if not category_dir.exists():
continue
@@ -622,7 +611,6 @@ def find_examples(category_indices: dict[str, Index], generate_main_index: bool)
examples.append(Example(path))
# Find examples in subdirectories
for path in EXAMPLE_DIR.glob("*/*.md"):
# Skip categorised examples
if path.parent.name in category_indices:
continue
examples.append(Example(path.parent))
@@ -643,7 +631,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic
if example.category not in ["training", "distillation"]:
continue
# Use mapped directory name if available
dir_name = category_dir_mapping.get(example.category, example.category)
category_dir = EXAMPLE_DIR / dir_name
relative_path = example.path.relative_to(category_dir)
@@ -656,7 +643,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic
model = path_parts[1] # e.g., "wan_i2v_14b_480p"
dataset = path_parts[2] # e.g., "crush_smol"
# Initialize nested structure
if example.category not in nested_structures:
nested_structures[example.category] = {}
if method not in nested_structures[example.category]:
@@ -664,7 +650,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic
if model not in nested_structures[example.category][method]:
nested_structures[example.category][method][model] = {}
# Store the nested structure
nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category,
method=method,
model=model,
@@ -677,7 +662,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic
dataset = path_parts[1] # e.g., "Wan-Syn-Data-480P"
method = "DMD" # Default method for distillation
# Initialize nested structure
if example.category not in nested_structures:
nested_structures[example.category] = {}
if method not in nested_structures[example.category]:
@@ -685,7 +669,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic
if model not in nested_structures[example.category][method]:
nested_structures[example.category][method][model] = {}
# Store the nested structure
nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category,
method=method,
model=model,
@@ -700,7 +683,7 @@ def generate_flat_examples(examples: list[Example], category_indices: dict[str,
"""Generate documentation for flat structure examples (inference, etc.)."""
for example in examples:
if example.category in ["training", "distillation"]:
continue # Skip nested structure examples
continue
# Determine which index to use for this example
if example.category is not None and example.category in category_indices:
@@ -711,7 +694,6 @@ def generate_flat_examples(examples: list[Example], category_indices: dict[str,
else:
continue
# Generate the example documentation
doc_path = index.path.parent / f"{example.path.stem}.md"
with open(doc_path, "w+", encoding="utf-8") as f:
f.write(example.generate())
@@ -729,7 +711,6 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[str, di
category_base_dir = category_index.path.parent
for method, models in nested_structures[category_name].items():
# Create method-level index
method_index = Index(path=category_base_dir / f"{method}.md",
title=fix_case(method),
description=f"Examples using {method}.",
@@ -737,35 +718,28 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[str, di
maxdepth=2)
for model, datasets in models.items():
# Generate dataset examples using the Example class
for dataset, nested_struct in datasets.items():
doc_path = category_base_dir / f"{nested_struct.filename}.md"
with open(doc_path, "w+", encoding="utf-8") as f:
f.write(nested_struct.example.generate())
# Create model-level index
model_index = Index(path=category_base_dir / f"{model}.md",
title=fix_case(model.replace('_', ' ')),
description=f"Examples for the {model} model.",
caption=f"{fix_case(model.replace('_', ' '))} Datasets",
maxdepth=1)
# Add dataset indices to model index
for dataset, nested_struct in datasets.items():
model_index.documents.append(nested_struct.filename)
# Write model index
with open(model_index.path, "w+", encoding="utf-8") as f:
f.write(model_index.generate())
# Add model to method index
method_index.documents.append(model)
# Write method index
with open(method_index.path, "w+", encoding="utf-8") as f:
f.write(method_index.generate())
# Add method to main category index
category_index.documents.append(method)
@@ -777,7 +751,6 @@ def generate_examples(generate_main_index: bool = False) -> None:
generate_main_index (bool): Whether to generate the main examples index.
If False, only category-specific indices will be generated.
"""
# Create category indices
category_indices = create_category_indices()
# Create the main examples index only if requested
@@ -796,29 +769,22 @@ def generate_examples(generate_main_index: bool = False) -> None:
# Find all examples
examples = find_examples(category_indices, generate_main_index)
# Create nested structures for training and distillation
nested_structures = create_nested_structures(examples)
# Generate flat structure examples (inference, etc.)
generate_flat_examples(examples, category_indices, examples_index, generate_main_index)
# Generate nested structure examples (training, distillation)
generate_nested_examples(nested_structures, category_indices)
# Generate the index files for categories
for category_index in category_indices.values():
if category_index.documents:
# Add to main index if it exists
if generate_main_index and examples_index:
main_index_dir = examples_index.path.parent
rel_path = os.path.relpath(category_index.path, start=main_index_dir)
examples_index.documents.insert(0, str(rel_path).replace("\\", "/").replace(".md", ""))
# Write the category index file
with open(category_index.path, "w+", encoding="utf-8") as f:
f.write(category_index.generate())
# Write the main index file if requested
if generate_main_index and examples_index:
with open(examples_index.path, "w+", encoding="utf-8") as f:
f.write(examples_index.generate())
examples/distill/MatrixGame2.0/distill_dmd.sh
+1 -10
View File
@@ -7,7 +7,6 @@ export TOKENIZERS_PARALLELISM=false
RUN_NAME=$(date +"%m%d_%H%M")
echo "RUN_NAME: $RUN_NAME"
# Model paths for Self-Forcing DMD distillation.
GENERATOR_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers"
REAL_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" # Teacher
FAKE_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" # Critic
@@ -16,7 +15,6 @@ DATA_DIR="data/matrixgame2"
VALIDATION_DATASET_FILE="examples/distill/MatrixGame2.0/validation.json"
NUM_GPUS=1
# Training arguments
training_args=(
--tracker_project_name "matrixgame2_sf"
--output_dir "checkpoints/matrixgame2_sf_${RUN_NAME}"
@@ -32,12 +30,8 @@ training_args=(
--simulate_generator_forward
--num_frames 81
--num_frame_per_block 3 # Frame generation block size for self-forcing
# --enable_gradient_masking
# --gradient_mask_last_n_frames 21
# --init_weights_from_safetensors "path/to/generator_ema.safetensors"
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size 1
@@ -47,7 +41,7 @@ parallel_args=(
)
model_args=(
--model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script
--model_path $GENERATOR_MODEL_PATH
--pretrained_model_name_or_path $GENERATOR_MODEL_PATH
--real_score_model_path $REAL_SCORE_MODEL_PATH
--fake_score_model_path $FAKE_SCORE_MODEL_PATH
@@ -58,7 +52,6 @@ dataset_args=(
--dataloader_num_workers 4
)
# Validation arguments
validation_args=(
--log_validation
--log_visualization
@@ -69,7 +62,6 @@ validation_args=(
--validation_guidance_scale "6.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 3e-6
--mixed_precision "bf16"
@@ -80,7 +72,6 @@ optimizer_args=(
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--checkpoints_total_limit 3
examples/distill/MatrixGame2.0/distill_dmd.slurm
-1
View File
@@ -77,7 +77,6 @@ training_args=(
--simulate_generator_forward
--num_frames 81
--num_frame_per_block 3
# --init_weights_from_safetensors "path/to/generator_ema.safetensors"
)
parallel_args=(
examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh
+1 -3
View File
@@ -33,8 +33,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_data_dir
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
training_args=(
--tracker_project_name SFwan_t2v_distill_self_forcing_dmd
@@ -64,7 +62,7 @@ parallel_args=(
)
model_args=(
--model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script
--model_path $GENERATOR_MODEL_PATH
--pretrained_model_name_or_path $GENERATOR_MODEL_PATH
--real_score_model_path $REAL_SCORE_MODEL_PATH
--fake_score_model_path $FAKE_SCORE_MODEL_PATH
examples/distill/SFWan2.2-A14B/distill_dmd.sh
-11
View File
@@ -23,7 +23,6 @@ nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
export MASTER_ADDR=${nodes[0]}
export TOKENIZERS_PARALLELISM=false
export WANDB_API_KEY=YOUR_WANDB_API_KEY
# export WANDB_API_KEY='your_wandb_api_key_here'
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
@@ -31,23 +30,13 @@ export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
# Configs
NUM_GPUS=8
# Model paths for Self-Forcing DMD distillation with Wan2.2:
# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2
# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
training_args=(
--tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd # Updated for Wan2.2
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
OUTPUT_DIR="checkpoints/wan_t2v_finetune"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
OUTPUT_DIR="checkpoints/wan_t2v_finetune"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
OUTPUT_DIR="checkpoints/wan_t2v_finetune"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh
-3
View File
@@ -32,7 +32,6 @@ export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
export WANDB_API_KEY=your_wandb_api_key
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -45,8 +44,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DIR=your_validation_path #(example:validation_64.json)
OUTPUT_DIR="checkpoints/wan_t2v_finetune"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh
-3
View File
@@ -32,7 +32,6 @@ export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
export WANDB_API_KEY=your_wandb_api_key
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -44,8 +43,6 @@ REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DIR=your_validation_path #(example:validation_64.json)
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh
-3
View File
@@ -9,7 +9,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
# Configs
NUM_GPUS=1
@@ -19,8 +18,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
OUTPUT_DIR="checkpoints/wan_t2v_finetune"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh
-3
View File
@@ -9,7 +9,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
# Configs
NUM_GPUS=1
@@ -18,8 +17,6 @@ REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/inference/basic/basic_matrixgame2_streaming.py
-1
View File
@@ -63,7 +63,6 @@ async def main():
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=True,
)
examples/inference/basic/basic_wan2_2.py
-1
View File
@@ -17,7 +17,6 @@ def main():
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=True,
)
examples/inference/eval/all_metrics_demo.py
+1 -1
View File
@@ -118,7 +118,7 @@ def score_all_metrics(video_path: str) -> None:
for name in ev.metric_names:
r = results[0].get(name)
if r is None:
continue # set metric, will show in corpus section
continue
if r.score is None:
print(f" {name}: SKIPPED — {r.details.get('skipped', 'no score')}")
else:
examples/inference/gradio/serving/ray_serve_backend.py
+2 -2
View File
@@ -187,8 +187,8 @@ class BaseModelDeployment:
num_gpus=1,
use_fsdp_inference=True,
text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V
dit_precision="fp32", # TODO: hardocde for I2V
dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
dit_precision="fp32",
dit_cpu_offload=config["dit_cpu_offload"],
vae_cpu_offload=config["vae_cpu_offload"],
VSA_sparsity=config["VSA_sparsity"],
examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v_vsa.yaml
-1
View File
@@ -35,7 +35,6 @@ method:
enable_gradient_in_rollout: true
start_gradient_frame: 0
# Critic optimizer
fake_score_learning_rate: 8.0e-6
fake_score_betas: [0.0, 0.999]
fake_score_lr_scheduler: constant
examples/train/configs/example.yaml
-12
View File
@@ -1,4 +1,3 @@
# ==============================================================================
# Full configuration reference for fastvideo.train
#
# Legend:
@@ -8,15 +7,12 @@
# Keys depend on the _target_ class constructor / method_config.
# [RESOLVED] — parsed by PipelineConfig.from_kwargs(); auto-populated from
# the model's config files. Only scalar overrides are useful.
# ==============================================================================
# ------------------------------------------------------------------------------
# models: [FREE]
#
# Each role is instantiated via _target_(*, training_config=..., **kwargs).
# Keys here are constructor kwargs of the _target_ class (e.g. WanModel).
# You can define any role name (student, teacher, critic, etc.).
# ------------------------------------------------------------------------------
models:
student:
_target_: fastvideo.train.models.wan.WanModel # required
@@ -39,13 +35,11 @@ models:
trainable: true
disable_custom_init_weights: true
# ------------------------------------------------------------------------------
# method: [FREE]
#
# Instantiated via _target_(*, cfg=RunConfig, role_models=...).
# All keys besides _target_ are available in self.method_config (a plain dict).
# Keys depend entirely on the method class.
# ------------------------------------------------------------------------------
method:
_target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method # required
@@ -71,11 +65,9 @@ method:
# --- FineTuneMethod keys (if using finetune instead) ---
# _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
# ------------------------------------------------------------------------------
# training: [TYPED] -> TrainingConfig
#
# Every field below has a typed default. Unknown keys are ignored.
# ------------------------------------------------------------------------------
training:
# --- training.distributed [TYPED] -> DistributedConfig ---
@@ -166,13 +158,11 @@ training:
dit_precision: fp32 # default: "fp32" (master weight precision)
# model_path: ... # default: "" (auto-derived from models.student.init_from)
# ------------------------------------------------------------------------------
# callbacks: [FREE]
#
# Each callback is instantiated via _target_(*, **kwargs).
# The callback name (e.g. "grad_clip") is arbitrary — only _target_ matters.
# training_config is injected automatically (not from YAML).
# ------------------------------------------------------------------------------
callbacks:
# --- GradNormClipCallback ---
@@ -199,13 +189,11 @@ callbacks:
output_dir: null # default: null (falls back to training.checkpoint.output_dir)
sampling_timesteps: null # default: null (explicit timestep list for streaming)
# ------------------------------------------------------------------------------
# pipeline: [RESOLVED] -> PipelineConfig
#
# Parsed by PipelineConfig.from_kwargs(). Most fields are auto-populated from
# the model's config files (vae_config, dit_config, text_encoder_configs, etc.).
# Only scalar overrides are typically needed here.
# ------------------------------------------------------------------------------
pipeline:
flow_shift: 3 # default: null (model-specific)
# LTX-2 linear QAT (SM100+ with FlashInfer); distinct from the role-local
examples/train/configs/knowledge_distillation/wan/causal_t2v.yaml
+1 -1
View File
@@ -24,7 +24,7 @@ models:
_target_: fastvideo.train.models.wan.WanCausalModel
init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
trainable: true
teacher: # remove once cache is complete
teacher:
_target_: fastvideo.train.models.wan.WanModel
init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
trainable: false
examples/train/configs/knowledge_distillation/wan/t2v.yaml
+1 -1
View File
@@ -21,7 +21,7 @@ models:
_target_: fastvideo.train.models.wan.WanModel
init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
trainable: true
teacher: # remove once cache is complete
teacher:
_target_: fastvideo.train.models.wan.WanModel
init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
trainable: false
examples/train/run.sh
-3
View File
@@ -27,7 +27,6 @@ if [[ "${CONFIG}" != /* ]]; then
CONFIG="$(pwd)/${CONFIG}"
fi
# ── GPU / node settings ──────────────────────────────────────────
NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}"
NUM_GPUS="${NUM_GPUS:-1}"
NNODES="${NNODES:-1}"
@@ -38,11 +37,9 @@ export TOKENIZERS_PARALLELISM=false
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
cd "${REPO_ROOT}"
# ── W&B ──────────────────────────────────────────────────────────
export WANDB_API_KEY="${WANDB_API_KEY:-}"
export WANDB_MODE="${WANDB_MODE:-online}"
# ── Log file ─────────────────────────────────────────────────────
CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
LOG_DIR="${LOG_DIR:-examples/train}"
examples/train/run_slurm.sh
-2
View File
@@ -49,7 +49,6 @@ JOB_NAME="${JOB_NAME:-${CONFIG_NAME}}"
OUTPUT_DIR="${OUTPUT_DIR:-logs/slurm}"
mkdir -p "${OUTPUT_DIR}"
# ── Build sbatch args ─────────────────────────────────────────────
# Forward the exported W&B API key through Slurm's environment so the generated
# job script and its xtrace output do not contain the credential value.
SBATCH_ARGS=(
@@ -70,7 +69,6 @@ SBATCH_ARGS=(
if [[ -n "${EXCLUDE}" ]]; then
SBATCH_ARGS+=(--exclude="${EXCLUDE}")
fi
# Apply hardware and wall-clock requirements selected by an experiment launcher.
if [[ -n "${SBATCH_CONSTRAINT}" ]]; then
SBATCH_ARGS+=(--constraint="${SBATCH_CONSTRAINT}")
fi
examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh
-1
View File
@@ -1,5 +1,4 @@
#!/usr/bin/env bash
# Run the modular Attn-QAT finetune recipe.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml
-3
View File
@@ -45,9 +45,6 @@ training:
checkpoints_total_limit: 3
resume_from_checkpoint: latest
# tracker:
# project_name: fastvideo
# run_name: matrixgame2_i2v_finetune
tracker:
trackers: [none]
project_name: ""
examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml
-1
View File
@@ -33,7 +33,6 @@ method:
enable_gradient_in_rollout: true
start_gradient_frame: 0
# Critic optimizer
fake_score_learning_rate: 3.0e-7
fake_score_betas: [0.9, 0.95]
fake_score_lr_scheduler: constant
examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml
-1
View File
@@ -6,7 +6,6 @@ models:
_target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel
init_from: mignonjia/mg_sf_distilled_zelda_1k_steps
trainable: true
# transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors
teacher:
_target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model
init_from: mignonjia/mg_bidirectional_zelda
examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh
-1
View File
@@ -8,7 +8,6 @@ MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=1
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm
-1
View File
@@ -29,7 +29,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_API_KEY=your-wandb-api-key
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
examples/training/finetune/MatrixGame2.0/finetune_df.sh
-4
View File
@@ -3,14 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers"
DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
@@ -80,7 +77,6 @@ miscellaneous_args=(
--ema_start_step 0
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
examples/training/finetune/MatrixGame2.0/finetune_i2v.sh
-3
View File
@@ -3,14 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers"
DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh
-1
View File
@@ -5,7 +5,6 @@ MODEL_PATH="Matrix-Game-2.0-Base-Diffusers"
DATA_MERGE_PATH="footsies-dataset/merge.txt"
OUTPUT_DIR="footsies-dataset/preprocessed/"
# export CUDA_VISIBLE_DEVICES=0
export MASTER_ADDR=localhost
export MASTER_PORT=29500
export RANK=0
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh
-4
View File
@@ -3,14 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
@@ -80,7 +77,6 @@ miscellaneous_args=(
--enable_gradient_checkpointing_type "full"
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm
-3
View File
@@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -41,8 +40,6 @@ NUM_GPUS=8
MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -41,8 +40,6 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Core training arguments
core_training_args=(
examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -41,8 +40,6 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Core training arguments
core_training_args=(
examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh
+1 -12
View File
@@ -1,6 +1,5 @@
# Basic info
export WANDB_MODE="online"
export NCCL_P2P_DISABLE=1
export TORCH_NCCL_ENABLE_MONITORING=0
@@ -11,15 +10,13 @@ export NODE_RANK=0
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
export WANDB_API_KEY="your_wandb_api_key_here" # TODO: Replace with your actual key or load from a secure location
export WANDB_API_KEY="your_wandb_api_key_here"
# Configs
NUM_GPUS=4
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR=data/Wan-Syn_77x448x832_600k
VALIDATION_DATASET_FILE=examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
# Core training arguments
core_training_args=(
--tracker_project_name wan_t2v_VSA
--output_dir "checkpoints/wan_t2v_finetune_VSA"
@@ -31,14 +28,12 @@ core_training_args=(
--enable_gradient_checkpointing_type "full" # if OOM enable this
)
# Validation generation shape arguments (used during training validation)
validation_generation_shape_args=(
--num_height 448
--num_width 832
--num_frames 77
)
# Parallel arguments
parallel_args=(
--num_gpus $NUM_GPUS
--sp_size 1
@@ -47,19 +42,16 @@ parallel_args=(
--hsdp_shard_dim 1
)
# Model arguments
model_args=(
--model_path "$MODEL_PATH"
--pretrained_model_name_or_path "$MODEL_PATH"
)
# Dataset arguments
dataset_args=(
--data_path "$DATA_DIR"
--dataloader_num_workers 4
)
# Validation arguments
validation_args=(
--log_validation
--validation_dataset_file "$VALIDATION_DATASET_FILE"
@@ -68,7 +60,6 @@ validation_args=(
--validation_guidance_scale "5.0"
)
# Optimizer arguments
optimizer_args=(
--learning_rate 1e-6
--mixed_precision "bf16"
@@ -78,7 +69,6 @@ optimizer_args=(
--max_grad_norm 1.0
)
# Miscellaneous arguments
miscellaneous_args=(
--inference_mode False
--checkpoints_total_limit 3
@@ -89,7 +79,6 @@ miscellaneous_args=(
--seed 1000
)
# VSA arguments
vsa_args=(
--VSA_decay_rate 0.03
--VSA_decay_interval_steps 50
examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm
-3
View File
@@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -41,8 +40,6 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
DATA_DIR=your_data_dir
VALIDATION_DATASET_FILE=your_validation_dataset_file
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Core training arguments
core_training_args=(
examples/training/finetune/cosmos2_5/finetune_t2w.sh
-1
View File
@@ -16,7 +16,6 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="KyleShao/Cosmos-Predict2.5-2B-Diffusers"
DATA_DIR="data/cosmos2_5_processed_t2w/combined_parquet_dataset/"
examples/training/finetune/cosmos2_5/finetune_t2w_lora.sh
-1
View File
@@ -10,7 +10,6 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="KyleShao/Cosmos-Predict2.5-2B-Diffusers"
DATA_DIR="data/cosmos2_5_processed_t2w/combined_parquet_dataset/"
examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh
-2
View File
@@ -7,9 +7,7 @@ set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
# ---------------------------------------------------------------------------
# Edit the DDP run here. Do not launch while another job owns the same GPUs.
# ---------------------------------------------------------------------------
FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train"
VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val"
EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth"
examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh
-2
View File
@@ -6,9 +6,7 @@ set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
# ---------------------------------------------------------------------------
# Edit the training run here. No environment variables are required.
# ---------------------------------------------------------------------------
VARIANT="small_44k"
FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train"
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh
-4
View File
@@ -3,14 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=8
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
@@ -80,7 +77,6 @@ miscellaneous_args=(
--enable_gradient_checkpointing_type "full"
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm
-3
View File
@@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -41,8 +40,6 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh
-4
View File
@@ -2,14 +2,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
@@ -79,7 +76,6 @@ miscellaneous_args=(
--ema_start_step 0
)
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
-2
View File
@@ -3,13 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
# Training arguments
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm
-3
View File
@@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
export TOKENIZERS_PARALLELISM=false
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
echo "MASTER_ADDR: $MASTER_ADDR"
echo "NODE_RANK: $NODE_RANK"
@@ -40,8 +39,6 @@ NUM_GPUS=8
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
# Training arguments
training_args=(
examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh
-2
View File
@@ -2,13 +2,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=1
# export CUDA_VISIBLE_DEVICES=4,5
# Training arguments
examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh
-2
View File
@@ -3,13 +3,11 @@
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
NUM_GPUS=4
# export CUDA_VISIBLE_DEVICES=4,5
mkdir -p ../profiler_traces/wan_t2v_finetune/
fastvideo-kernel/attn_qat_infer/blackwell/api.cu
-19
View File
@@ -68,9 +68,7 @@ void set_params_fprop(Flash_fwd_params &params,
bool single_level_p_quant=false,
bool seqlenq_ngroups_swapped=false) {
// Reset the parameters
params = {};
// Set the pointers and strides.
params.q_ptr = q.data_ptr();
params.k_ptr = k.data_ptr();
params.v_ptr = v.data_ptr();
@@ -119,13 +117,11 @@ void set_params_fprop(Flash_fwd_params &params,
params.cu_seqlens_k = static_cast<int *>(cu_seqlens_k_d);
params.seqused_k = static_cast<int *>(seqused_k);
// P = softmax(QK^T)
params.p_ptr = p_d;
// Softmax sum
params.softmax_lse_ptr = softmax_lse_d;
// Set the dimensions.
params.b = b;
params.h = h;
params.h_k = h_k;
@@ -140,7 +136,6 @@ void set_params_fprop(Flash_fwd_params &params,
params.head_divmod = cutlass::FastDivmod(int(h));
// Set the different scale values.
params.scale_softmax = softmax_scale;
params.scale_softmax_log2 = softmax_scale * M_LOG2E;
__half scale_softmax_log2_half = __float2half(params.scale_softmax_log2);
@@ -149,10 +144,6 @@ void set_params_fprop(Flash_fwd_params &params,
// Set this to probability of keeping an element to simplify things.
params.p_dropout = 1.f - p_dropout;
// Convert p from float to int so we don't have to convert the random uint to float to compare.
// [Minor] We want to round down since when we do the comparison we use <= instead of <
// params.p_dropout_in_uint = uint32_t(std::floor(params.p_dropout * 4294967295.0));
// params.p_dropout_in_uint16_t = uint16_t(std::floor(params.p_dropout * 65535.0));
params.p_dropout_in_uint8_t = uint8_t(std::floor(params.p_dropout * 255.0));
params.rp_dropout = 1.f / params.p_dropout;
params.scale_softmax_rp_dropout = params.rp_dropout * params.scale_softmax;
@@ -276,18 +267,12 @@ mha_fwd_impl(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head
CHECK_SHAPE(q, batch_size, num_heads, seqlen_q, head_size_og);
CHECK_SHAPE(k, batch_size, num_heads_k, seqlen_k, head_size_og);
CHECK_SHAPE(v, batch_size, num_heads_k, unpacked_head_size, seqlen_k/2);
// CHECK_SHAPE(delta_s, batch_size, num_heads, seqlen_q / 128, seqlen_k);
// CHECK_SHAPE(sfq, batch_size, seqlen_q, num_heads, unpacked_head_size);
// CHECK_SHAPE(sfk, batch_size, seqlen_k, num_heads_k, unpacked_head_size);
// CHECK_SHAPE(sfv, batch_size, unpacked_head_size, num_heads_k, seqlen_k);
TORCH_CHECK(unpacked_head_size % 8 == 0, "head_size must be a multiple of 8");
auto dtype = is_bf16 ? at::ScalarType::BFloat16 : at::ScalarType::Half;
at::Tensor out = torch::empty({batch_size, num_heads, seqlen_q, unpacked_head_size}, opts.dtype(dtype));
auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
// const int head_size = round_multiple(head_size_og, 8);
// const int head_size_rounded = round_multiple(head_size, 32);
const int seqlen_q_rounded = round_multiple(seqlen_q, flash::BLOCK_M);
const int seqlen_k_rounded = round_multiple(seqlen_k, flash::BLOCK_N);
@@ -348,10 +333,6 @@ mha_fwd_impl(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head
// if (out_.has_value()) { out_.value().copy_(out); }
// }
// return {out, q_padded, k_padded, v_padded, out_padded, softmax_lse, p};
// cudaDeviceSynchronize();
// auto err = cudaGetLastError();
// printf("%s\n", cudaGetErrorString(err));
return {out, softmax_lse};
}
fastvideo-kernel/attn_qat_infer/blackwell/block_info.h
-2
View File
@@ -20,7 +20,6 @@
namespace flash {
////////////////////////////////////////////////////////////////////////////////////////////////////
template<bool Varlen=true>
struct BlockInfo {
@@ -55,6 +54,5 @@ struct BlockInfo {
const int actual_seqlen_k;
};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace flash
fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h
-4
View File
@@ -119,7 +119,6 @@ struct CollectiveEpilogueFwd{
auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
constexpr int numel = decltype(size(tOrO))::value;
cutlass::NumericArrayConverter<Element, float, numel> convert_op;
// HACK: this requires tensor to be "contiguous"
auto frag = convert_op(*reinterpret_cast<const cutlass::Array<float, numel> *>(tOrO.data()));
auto tOrO_out = make_tensor(make_rmem_ptr<Element>(&frag), tOrO.layout());
Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out); // ((Atom,AtomNum), MMA_M, MMA_N)
@@ -183,7 +182,6 @@ struct CollectiveEpilogueFwd{
tma_store_wait<0>();
}
// Write 0 to output and -inf to LSE
CUTLASS_DEVICE void
store_zero(
Params const& epilogue_params,
@@ -202,9 +200,7 @@ struct CollectiveEpilogueFwd{
Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
Tensor tOrO = make_fragment_like(tOgO);
clear(tOrO);
// Construct identity layout for sO
Tensor cO = cute::make_identity_tensor(select<0, 2>(TileShape_MNK{})); // (BLK_M,BLK_K) -> (blk_m,blk_k)
// Repeat the partitioning with identity layouts
Tensor tOcO = gmem_thr_copy_O.partition_D(cO);
Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOgO)));
#pragma unroll
fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h
+1 -4
View File
@@ -97,7 +97,6 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp,
CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params);
}
// Obtain warp index
PipelineParams pipeline_params_v;
pipeline_params_v.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV;
@@ -142,7 +141,7 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp,
cutlass::arch::warpgroup_reg_dealloc<24>();
TileScheduler scheduler;
if (producer_warp_role == ProducerWarpRole::Mainloop) { // Load Q, K, V
if (producer_warp_role == ProducerWarpRole::Mainloop) {
PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
PipelineState smem_pipe_write_k = cutlass::make_producer_start_state<MainloopPipeline>();
PipelineState smem_pipe_write_v = cutlass::make_producer_start_state<MainloopPipeline>();
@@ -179,8 +178,6 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp,
for (auto work_tile_info = scheduler.get_initial_work(); work_tile_info.is_valid(scheduler_params); work_tile_info = scheduler.get_next_work(scheduler_params, work_tile_info)) {
// Attention output (GEMM-II) accumulator.
Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 2>(TileShape_MNK{}));
// flash::Softmax<2 * (2 * kBlockM / NumMmaThreads)> softmax;
// Pass single_level_p_quant flag to control P quantization mode
flash::SoftmaxFused<2 * (2 * kBlockM / NumMmaThreads)> softmax_fused(params.single_level_p_quant);
auto block_coord = work_tile_info.get_block_coord(scheduler_params);
auto [m_block, bidh, bidb] = block_coord;
fastvideo-kernel/attn_qat_infer/blackwell/launch.h
-1
View File
@@ -79,7 +79,6 @@ void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) * size<0>(ClusterShape{});
typename Scheduler::Arguments scheduler_args = {num_blocks_m, params.h, params.b};
typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args);
// Get the ptr to kernel function.
void *kernel;
kernel = (void *)flash::compute_attn_ws<Kernel_traits, Is_causal, Scheduler>;
int smem_size = sizeof(typename Kernel_traits::SharedStorage);
fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h
-1
View File
@@ -683,7 +683,6 @@ struct CollectiveMainloopFwd {
Tensor tOrSFP = make_tensor<ElementSF>(LayoutSFP{});
Tensor tOrSFP_flt = filter_zeros(tOrSFP);
Tensor tSrDS = make_tensor<float>(make_shape(_8{}, _4{}), make_stride(_1{}, _8{}));
// copy qk and sf from smem to rmem
auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtomQ{}, tiled_mma_qk);
auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(thread_idx);
Tensor tSsQ = smem_thr_copy_Q.partition_S(as_position_independent_swizzle_tensor(sQ));
fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h
-1
View File
@@ -65,7 +65,6 @@ public:
// MSVC requires protecting use of CUDA-specific nonstandard syntax,
// like blockIdx and gridDim, with __CUDA_ARCH__.
#if defined(__CUDA_ARCH__)
// current_work_linear_idx_ = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
current_work_linear_idx_ = blockIdx.x;
#else
CUTLASS_ASSERT(false && "This line should never be reached");
fastvideo-kernel/attn_qat_infer/blackwell/utils.h
-2
View File
@@ -136,7 +136,6 @@ __forceinline__ __device__ __half2 half_exp(__half2 x) {
return out;
}
// Apply the exp to all the elements.
template <bool zero_init=false, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
__forceinline__ __device__ void max_scale_exp2_sum(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> &max, Tensor<Engine1, Layout1> &sum, const float scale) {
static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
@@ -164,7 +163,6 @@ __forceinline__ __device__ void max_scale_exp2_sum(Tensor<Engine0, Layout0> &ten
}
}
// Apply the exp to all the elements.
template <bool Scale_max=true, bool Check_inf=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
__forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> const &max, const float scale) {
static_assert(Layout0::rank == 2, "Only support 2D Tensor");
fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py
-1
View File
@@ -79,7 +79,6 @@ if n % 128 != 0:
else:
q_padded = q
# use torch transpose + scaled_fp4_quant to get the ground truth
q_padded = q_padded.transpose(2, 3).reshape(b, h, n_padded, d).contiguous()
o_gt = torch.empty((b, h, n_padded, d // 2), dtype=torch.uint8, device='cuda')
o_s_gt = torch.empty((b, h, n_padded, d // 16), dtype=torch.float8_e4m3fn, device='cuda')
fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py
-4
View File
@@ -127,11 +127,9 @@ def bench_kineto(fn,
if not using_nsys:
profiler.step()
# Return 1 if using Nsight Systems
if using_nsys:
return 1
# Parse the profiling table
assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple)
is_tupled = isinstance(kernel_names, tuple)
prof_lines = profiler.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n')
@@ -140,11 +138,9 @@ def bench_kineto(fn,
for name in kernel_names:
assert sum([name in line for line in prof_lines]) == 1, f'Errors of the kernel {name} in the profiling table'
# Save chrome traces
if trace_path is not None:
profiler.export_chrome_trace(trace_path)
# Return average kernel times
units = {'ms': 1e3, 'us': 1e6}
kernel_times = []
for name in kernel_names:
fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu
-11
View File
@@ -125,7 +125,6 @@
using Type = __nv_bfloat162;
};
// Define a 32 bytes packed data type.
template <class Type>
struct PackedVec {
typename TypeConverter<Type>::Type elts[8];
@@ -152,7 +151,6 @@
constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD;
// load input
const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN;
int load_token_id;
@@ -183,7 +181,6 @@
(threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; // feature dim
}
// calculate max of every consecutive 16 elements
auto localMax = __habs2(in_vec.elts[0]);
#pragma unroll
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { // local max
@@ -204,7 +201,6 @@
float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue;
// convert input to float2 and apply scale
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
#pragma unroll
@@ -218,7 +214,6 @@
fp2Vals[i].y = fp2Vals[i].y * SFValueInv;
}
// convert to e2m1
uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8];
#pragma unroll
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) {
@@ -229,7 +224,6 @@
// is not a multiple of BLOCK_SIZE (out-of-bounds write fix).
if (token_id >= num_tokens) return;
// save
if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) {
reinterpret_cast<uint32_t*>(output +
batch_id * stride_bz_output +
@@ -284,7 +278,6 @@
constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD;
constexpr uint32_t NUM_THREADS_PER_SEQ = BLOCK_SIZE / CVT_FP4_ELTS_PER_THREAD;
// load input
const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN;
PackedVec in_vec;
@@ -312,7 +305,6 @@
in_vec.elts[i].y = shared_input[(threadIdx.x / NUM_THREADS_PER_SEQ) + ((threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD + 2 * i + 1) * head_dim];
}
// calculate max of every consecutive 16 elements
auto localMax = __habs2(in_vec.elts[0]);
#pragma unroll
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { // local max
@@ -333,7 +325,6 @@
float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue;
// convert input to float2 and apply scale
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
#pragma unroll
@@ -347,7 +338,6 @@
fp2Vals[i].y = fp2Vals[i].y * SFValueInv;
}
// convert to e2m1
uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8];
#pragma unroll
for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) {
@@ -360,7 +350,6 @@
(threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD;
if (write_token_id >= num_tokens) return;
// save
if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) {
reinterpret_cast<uint32_t*>(output +
batch_id * stride_bz_output +
fastvideo-kernel/build.sh
-2
View File
@@ -224,8 +224,6 @@ export CMAKE_ARGS
echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-<unset>}"
echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
# Build and install
# Use -v for verbose output
if [ -n "${WHEEL_DIR}" ]; then
mkdir -p "${WHEEL_DIR}"
uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" .
fastvideo-kernel/csrc/attention/block_sparse_h100.cu
+1 -32
View File
@@ -65,7 +65,7 @@ template<int D> struct fwd_globals {
template<int D>
__global__ __launch_bounds__(128, 4)
void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block size of 64
void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
#if FASTVIDEO_TK_HOPPER
extern __shared__ int __shm[];
tma_swizzle_allocator al((int*)&__shm[0]);
@@ -138,7 +138,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
// wait k
wait(k_smem_arrived, kv_idx % 2);
// compute QK^T
warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]);
copy(max_vec_last_scaled, max_vec);
@@ -147,7 +146,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
warpgroup::mma_async_wait();
// load K
if (threadIdx.x == 0) {
tma::expect_bytes(k_smem_arrived, sizeof(k_tile));
coord<k_tile> k_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0};
@@ -180,11 +178,9 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
// wait v
wait(v_smem_arrived, kv_idx % 2);
// compute SV
warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]);
warpgroup::mma_async_wait();
// load V
if (threadIdx.x == 0) {
tma::expect_bytes(v_smem_arrived, sizeof(v_tile));
coord<v_tile> v_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0};
@@ -198,7 +194,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
// wait k
wait(k_smem_arrived, kv_idx % 2);
// compute QK^T
warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]);
copy(max_vec_last_scaled, max_vec);
@@ -234,7 +229,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
// wait v
wait(v_smem_arrived, kv_idx % 2);
// compute SV
warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]);
warpgroup::mma_async_wait();
}
@@ -550,7 +544,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
if constexpr (D == 64) { mul(ds_block_t, ds_block_t, 0.125f); }
else { mul(ds_block_t, ds_block_t, 0.08838834764f); }
// load vec
if (threadIdx.x == 0) {
coord<l_tile> vec_idx = {blockIdx.z, blockIdx.y, 0, load_q_block_index};
tma::expect_bytes(vec_b[0], sizeof(l_smem[0]) + sizeof(d_smem[0]));
@@ -564,7 +557,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
warpgroup::store(ds_smem_t[0], ds_block_t);
warpgroup::mma_async_wait();
// load og
if (threadIdx.x == 0) {
coord<q_tile> tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0};
tma::expect_bytes(o_b[0], sizeof(og_smem[0]));
@@ -575,7 +567,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
warpgroup::mma_commit_group();
warpgroup::mma_async_wait();
// load q
if (threadIdx.x == 0) {
coord<q_tile> q_tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0};
tma::expect_bytes(q_b[0], sizeof(q_smem[0]));
@@ -587,11 +578,9 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
warpgroup::mma_commit_group();
warpgroup::mma_async_wait();
// store qg to shared memory
warpgroup::store(qg_smem, qg_reg);
__syncthreads();
// store and add dQ to global memory
if (threadIdx.x / 32 == 0) {
coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
tma::store_add_async(g.qg, qg_smem, tile_idx);
@@ -643,11 +632,9 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
warpgroup::mma_commit_group();
warpgroup::mma_async_wait();
// store qg to shared memory
warpgroup::store(qg_smem, qg_reg);
__syncthreads();
// store and add dQ to global memory
if (threadIdx.x / 32 == 0) {
coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
tma::store_add_async(g.qg, qg_smem, tile_idx);
@@ -655,7 +642,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
}
}
// store kq and vq
// ! the following two line seems unnecessary.
// tma::store_async_wait(); // ensure qg is finished
@@ -768,7 +754,6 @@ block_sparse_attention_forward(
float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
float* d_l = reinterpret_cast<float*>(l_ptr);
//cudadevicesynchronize();
const c10::cuda::OptionalCUDAGuard device_guard(q.device());
const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
@@ -820,7 +805,6 @@ block_sparse_attention_forward(
fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
CHECK_CUDA_ERROR(cudaGetLastError());
// cudaStreamSynchronize(stream);
}
if (head_dim == 128) {
@@ -871,11 +855,9 @@ block_sparse_attention_forward(
fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
CHECK_CUDA_ERROR(cudaGetLastError());
// cudaStreamSynchronize(stream);
}
return {o, l_vec};
//cudadevicesynchronize();
}
std::vector<torch::Tensor>
@@ -996,11 +978,9 @@ block_sparse_attention_backward(torch::Tensor q,
constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
int threads = PREP_NUM_WARPS * kittens::WARP_THREADS;
//cudadevicesynchronize();
const c10::cuda::OptionalCUDAGuard device_guard(q.device());
const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
// cudaStreamSynchronize(stream);
// TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
dim3 grid_bwd(q_seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
@@ -1084,7 +1064,6 @@ block_sparse_attention_backward(torch::Tensor q,
dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
threads = 128;
//cudadevicesynchronize();
{
cudaFuncSetAttribute(
@@ -1101,12 +1080,6 @@ block_sparse_attention_backward(torch::Tensor q,
bwd_attend_ker<64><<<grid_bwd_2, threads, 72000, stream>>>(bwd_global);
}
// CHECK_CUDA_ERROR(cudaGetLastError());
// cudaStreamSynchronize(stream);
//cudadevicesynchronize();
// const auto kernel_end = std::chrono::high_resolution_clock::now();
// std::cout << "Kernel Time: " << std::chrono::duration_cast<std::chrono::microseconds>(kernel_end - start).count() << "us" << std::endl;
// std::cout << "---" << std::endl;
}
if (head_dim == 128) {
@@ -1188,7 +1161,6 @@ block_sparse_attention_backward(torch::Tensor q,
dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
threads = 128;
//cudadevicesynchronize();
{
cudaFuncSetAttribute(
@@ -1200,10 +1172,7 @@ block_sparse_attention_backward(torch::Tensor q,
bwd_attend_ker<128><<<grid_bwd_2, threads, 113000, stream>>>(bwd_global);
}
// cudaStreamSynchronize(stream);
//cudadevicesynchronize();
}
return {qg, kg, vg};
//cudadevicesynchronize();
}
fastvideo-kernel/csrc/attention/st_attn_h100.cu
-5
View File
@@ -310,9 +310,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); }
warpgroup::mma_async_wait();
// apply non-pad mask
int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height;
// printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset);
right_fill(att_block, att_block, offset, base_types::constants<float>::neg_infty());
@@ -437,7 +435,6 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
float* d_l = reinterpret_cast<float*>(l_ptr);
//cudadevicesynchronize();
const c10::cuda::OptionalCUDAGuard device_guard(q.device());
const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
@@ -578,10 +575,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
}
CHECK_CUDA_ERROR(cudaGetLastError());
// cudaStreamSynchronize(stream);
}
return o;
//cudadevicesynchronize();
}
fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu
-1
View File
@@ -267,7 +267,6 @@ void ulysses_a2a_tuned(int64_t handle, torch::Tensor inp, torch::Tensor out, int
#undef LAUNCH_ULYSSES_A2A
TORCH_CHECK(cudaGetLastError() == cudaSuccess, "ulysses_a2a kernel launch failed");
// Copy this rank's completed result out of the window.
auto status = cudaMemcpyAsync(out.data_ptr(), ctx->buf, out_bytes, cudaMemcpyDeviceToDevice,
stream);
TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status));
fastvideo-kernel/csrc/common_extension.cpp
-1
View File
@@ -2,7 +2,6 @@
#include <ATen/ATen.h>
#include <vector>
// Forward declarations
#ifdef TK_COMPILE_ST_ATTN
extern torch::Tensor sta_forward(
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp
-6
View File
@@ -295,7 +295,6 @@ struct GemmKernel {
SmemLayoutB{}
);
// register
TiledMma tiled_mma;
auto thr_mma = tiled_mma.get_slice(t_idx);
auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));
@@ -314,7 +313,6 @@ struct GemmKernel {
clear(tDrD);
// global to shared copy
G2STiledCopyA g2s_tiled_copy_a;
auto g2s_thr_copy_a = g2s_tiled_copy_a.get_slice(t_idx);
auto tAgA = g2s_thr_copy_a.partition_S(gA);
@@ -371,7 +369,6 @@ struct GemmKernel {
int32_t fallbackA_load = 0;
int32_t fallbackA_mma = 0;
// copy first Stage - 1 tile
CUTLASS_PRAGMA_UNROLL
for (int i = 0, _i = min(Stage - 1, ntile); i < _i; ++i) {
if (g2s_b_tile < ntile) {
@@ -410,8 +407,6 @@ struct GemmKernel {
// a s2r increase anyway
s2r_a_smem = (s2r_a_smem + 1) % Stage;
// get next s2r b tile int64_t
// end of a block
// dequant first
dequant<AccumulatorDtype, TileM * TileN / ThreadNum, FastInt2Float>(
@@ -425,7 +420,6 @@ struct GemmKernel {
mma_block_a = mma_block_b;
if (mma_block_a < size<1>(AS)) scale_a = AS(blk_m_coord, mma_block_a);
// load next stage
if (g2s_b_tile < ntile) {
g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock;
copy_AB<IsEvenM>(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit);
fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp
-4
View File
@@ -61,7 +61,6 @@ public:
int tidx = threadIdx.x;
float x[NumElementPerThread];
// load
Loader<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> loader;
loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx);
@@ -86,7 +85,6 @@ public:
}
__syncthreads();
// var reduction
float v = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps);
CUTLASS_PRAGMA_UNROLL
@@ -94,7 +92,6 @@ public:
x[i] /= v;
if constexpr (Affine) {
// load weight
Loader<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> weight_loader;
float w[NumElementPerThread];
weight_loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx);
@@ -111,7 +108,6 @@ public:
}
}
// save y
{
Saver<OutputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven, false, false> saver;
if constexpr (std::is_same_v<OutputDtype, float>) {
fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp
-3
View File
@@ -56,14 +56,12 @@ public:
int tidx = threadIdx.x;
float x[NumElementPerThread];
// load
Loader<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> loader;
loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx);
// rms reduction
float rms = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps);
// load weight
Loader<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> weight_loader;
float w[NumElementPerThread];
loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx);
@@ -73,7 +71,6 @@ public:
for (int i = 0; i < NumElementPerThread; ++i)
x[i] = w[i] * x[i] / rms ;
// save y
OutputDtype *output_reg = (OutputDtype*)x;
if constexpr (!std::is_same_v<OutputDtype, float>) {
output_reg = (OutputDtype*)w;
fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp
-1
View File
@@ -109,7 +109,6 @@ public:
float float_reg[NumElementPerThread];
// load float32 data
Loader<InputDtype, BlockSize, BlockSize, NumThrPerCta, IsEvenM, IsEvenN> loader;
loader.load(params.Iptr, float_reg, params.m, params.n, blk_m, blk_n, tidx);
quantization(
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py
-27
View File
@@ -9,7 +9,6 @@ import triton.language as tl
from triton.tools.tensor_descriptor import TensorDescriptor
from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize
# from attn_qat_infer.api import triton_group_mean
def is_cuda():
@@ -141,7 +140,6 @@ def _attn_fwd_inner(acc,
elif STAGE == 2:
lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
lo = tl.multiple_of(lo, BLOCK_M)
# causal = False
else:
lo, hi = 0, N_CTX
offsetk_y = offset_y + lo # offset from the start of the current batch-head combination
@@ -149,7 +147,6 @@ def _attn_fwd_inner(acc,
offsetv_y = offset_y * HEAD_DIM + lo
else:
offsetv_y = offset_y + lo
# loop over k, v and update accumulator
for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize):
start_n = tl.multiple_of(start_n, BLOCK_N)
kv_valid = (start_n + offs_n) < N_CTX
@@ -191,7 +188,6 @@ def _attn_fwd_inner(acc,
acc = acc * alpha[:, None]
if IS_QAT:
high_prec_acc = high_prec_acc * alpha[:, None]
# prepare p and v for the dot
if dtype == tl.float8e5:
v = desc_v.load([0, offsetv_y]).T
else:
@@ -298,7 +294,6 @@ configs = [
for s in NUM_STAGES_OPTIONS for w in [4, 8]
]
if "PYTEST_VERSION" in os.environ:
# Use a single config in testing for reproducibility
configs = [
triton.Config(dict(BLOCK_M=128, BLOCK_N=64), num_stages=2, num_warps=4, pre_hook=_host_descriptor_pre_hook),
]
@@ -329,8 +324,6 @@ def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape):
return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape)
# @triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX_Q", "N_CTX_KV", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"],
# prune_configs_by={'early_config_prune': prune_invalid_configs}, cache_results=True)
@triton.jit
def _attn_fwd(
sm_scale,
@@ -398,11 +391,9 @@ def _attn_fwd(
offset_y_q = off_z * (N_CTX_Q * H) + off_h * N_CTX_Q # offset for query tensor
offset_y_kv = off_z * (N_CTX_KV * H) + off_h * N_CTX_KV # offset for key/value tensors
qo_offset_y = offset_y_q + start_m * BLOCK_M
# initialize offsets
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N) # for kv blocks
q_valid = offs_m < N_CTX_Q
# initialize pointer to m and l
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
@@ -411,10 +402,8 @@ def _attn_fwd(
else:
# dummy value
high_prec_acc = acc
# load scales
qk_scale = sm_scale
qk_scale *= 1.44269504 # 1/log(2)
# load q: it will stay in SRAM throughout
q = desc_q.load([qo_offset_y, 0]) # load from start of q block and start of the entire head dimension
q = tl.where(q_valid[:, None], q, 0.0)
# stage 1: off-band
@@ -434,7 +423,6 @@ def _attn_fwd(
BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV,
warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P,
two_level_quant_P, use_global_sf_P, JOIN_QAT_PV)
# stage 2: on-band
if STAGE & 2:
acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v,
offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM,
@@ -505,12 +493,10 @@ def _attn_bwd_preprocess(O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, HEAD_
off_hz = tl.program_id(1)
off_n = tl.arange(0, HEAD_DIM)
valid = off_m < N_CTX
# load
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None])
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :],
mask=valid[:, None]).to(tl.float32)
delta = tl.sum(o * do, axis=1)
# write-back
tl.store(Delta + off_hz * N_CTX + off_m, delta, mask=valid)
@@ -573,7 +559,6 @@ def _attn_bwd_dkdv(
qk = tl.where(mask, qk, -1.0e6)
p = tl.math.exp2(qk - m[:, None])
do = tl.load(do_ptrs, mask=q_valid[:, None])
# Compute dV.
p_quant = p
if IS_QAT and fake_quant_P:
p_quant, _ = fake_quantize(src_tensor=p,
@@ -586,7 +571,6 @@ def _attn_bwd_dkdv(
dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do)
# D (= delta) is pre-divided by ds_scale.
Di = tl.load(D + offs_m, mask=q_valid)
# Compute dP and dS.
dp = tl.dot(do, tl.trans(v))
ds = p * (dp - Di[:, None])
ds = ds.to(tl.bfloat16)
@@ -659,7 +643,6 @@ def _attn_bwd_dq(
mask = (offs_m[:, None] >= offs_n[None, :])
qk = tl.where(mask, qk, -1.0e6)
p = tl.math.exp2(qk - m)
# Compute dP and dS.
dp = tl.dot(do, tl.trans(v))
ds = p * (dp - Di[:, None])
# Compute dQ.
@@ -1037,7 +1020,6 @@ def _attn_bwd(
if SMOOTH_Q:
Q_MEAN += adj
# load scales
offs_k = tl.arange(0, HEAD_DIM)
start_n = pid * BLOCK_N1
@@ -1050,7 +1032,6 @@ def _attn_bwd(
dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
# load K and V: they stay in SRAM throughout the inner loop.
k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None])
@@ -1101,7 +1082,6 @@ def _attn_bwd(
start_m = 0
num_steps = (N_CTX + BLOCK_M1 - 1) // BLOCK_M1
# Compute dK and dV for non-masked blocks.
dk, dv = _attn_bwd_dkdv(dk,
dv,
Q,
@@ -1133,7 +1113,6 @@ def _attn_bwd(
dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
tl.store(dv_ptrs, dv, mask=kv_valid[:, None])
# Write back dK.
dk *= sm_scale
dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
tl.store(dk_ptrs, dk, mask=kv_valid[:, None])
@@ -1185,14 +1164,12 @@ def _attn_bwd(
warp_specialize=warp_specialize,
)
end_n -= num_steps * MASK_BLOCK_N2
# stage 2: process KV blocks from 0 to end_n (before diagonal), ceiling division to include remainder
num_steps = (end_n + BLOCK_N2 - 1) // BLOCK_N2
start_n = 0
else:
# For non-causal, process all KV blocks from 0 to N_CTX (ceiling division to include remainder)
num_steps = (N_CTX + BLOCK_N2 - 1) // BLOCK_N2
start_n = 0
# stage 2
dq = _attn_bwd_dq(
dq,
q,
@@ -1299,7 +1276,6 @@ class _attention(torch.autograd.Function):
if IS_QAT:
high_prec_o = torch.empty_like(q)
else:
# Initialize to a dummy value
high_prec_o = o
stage = 3 if causal else 1
extra_kern_args = {}
@@ -1442,7 +1418,6 @@ class _attention(torch.autograd.Function):
use_global_sf=use_global_sf_QKV,
)
# Apply pre-hook to set block shapes on tensor descriptors
_host_descriptor_pre_hook({
"BLOCK_M": fwd_block_m,
"BLOCK_N": fwd_block_n,
@@ -1647,7 +1622,6 @@ class _attention(torch.autograd.Function):
num_stages=dkdv_launch[1],
)
elif N_CTX_Q == N_CTX_KV:
# Use existing kernel for self-attention (same sequence lengths)
grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD)
_attn_bwd[grid](q,
arg_k,
@@ -1684,7 +1658,6 @@ class _attention(torch.autograd.Function):
num_warps=NUM_WARPS,
num_stages=NUM_STAGES)
else:
# Use separate kernels for cross-attention (different sequence lengths)
grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD)
_attn_bwd_dq_cross[grid_dq](
q,
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py
-10
View File
@@ -183,11 +183,9 @@ def _attn_bwd_preprocess(
off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
off_hz = tl.program_id(1)
off_n = tl.arange(0, HEAD_DIM)
# load
o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
delta = tl.sum(o * do, axis=1)
# write-back
tl.store(Delta + off_hz * N_CTX + off_m, delta)
@@ -255,13 +253,11 @@ def _attn_bwd_dkdv(
pT = tl.where(mask[:, None], pT, 0.0)
do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
# Compute dV.
ppT = pT
ppT = ppT.to(tl.bfloat16)
dv += tl.dot(ppT, do)
# D (= delta) is pre-divided by ds_scale.
Di = tl.load(D + offs_m)
# Compute dP and dS.
dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
dsT = pT * (dpT - Di[None, :])
dsT = dsT.to(tl.bfloat16)
@@ -333,7 +329,6 @@ def _attn_bwd_dq(
offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
mask = offs_in_block < block_size
p = tl.where(mask[None, :], p, 0.0)
# Compute dP and dS.
dp = tl.dot(do, vT).to(tl.float32)
ds = p * (dp - Di[:, None])
ds = ds.to(tl.bfloat16)
@@ -392,7 +387,6 @@ def _attn_bwd(
M += off_chz
D += off_chz
# load scales
offs_k = tl.arange(0, HEAD_DIM)
start_n = pid * BLOCK_N1
@@ -403,7 +397,6 @@ def _attn_bwd(
dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
# load K and V: they stay in SRAM throughout the inner loop.
k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
@@ -438,7 +431,6 @@ def _attn_bwd(
dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
tl.store(dv_ptrs, dv)
# Write back dK.
dk *= sm_scale
dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
tl.store(dk_ptrs, dk)
@@ -481,7 +473,6 @@ def _attn_bwd(
end_n,
num_steps #
)
# Write back dQ.
dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
dq *= sm_scale
tl.store(dq_ptrs, dq)
@@ -557,7 +548,6 @@ def _attn_bwd_dkdv_kernel(
start_n = kv_blk * BLOCK_N1
offs_n = start_n + tl.arange(0, BLOCK_N1)
# load K and V: they stay in SRAM throughout the inner loop.
k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py
-2
View File
@@ -165,7 +165,6 @@ def _compute_dequant(
if dst_dtype == tl.float16:
dst_scale = dst_scale.to(tl.float16)
# Now upcast the tensor.
intermediate_dtype: tl.constexpr = tl.bfloat16 if dst_dtype == tl.float32 else dst_dtype
if cuda_capability_geq(10, 0):
assert is_fp4
@@ -215,7 +214,6 @@ def _compute_dequant(
# 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type
x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0)
x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1)
# 3) x is zero, do nothing
dst_tensor = tl.interleave(x0, x1).to(intermediate_dtype, bitcast=True)
dst_tensor = dst_tensor.to(dst_dtype)
fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py
-8
View File
@@ -78,12 +78,10 @@ def _attn_fwd_loop(
exp_scores = tl.math.exp2(scores - new_m[:, None])
current_l = tl.sum(exp_scores, axis=1)
# Update L <- L * exp(M - M') + L1, M <- M'
alpha = tl.math.exp2(m - new_m)
l = l * alpha + current_l
m = new_m
# Update O <- O * exp(M - M') + P @ V
acc = (acc * alpha[:, None] + tl.dot(exp_scores.to(v.type.element_ty), v))
return m, l, acc
@@ -286,8 +284,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,
total_tile_size = tile_t * tile_h * tile_w
# BLOCK_Q=128
# BLOCK_KV=128
BLOCK_DIM = head_dim
output = torch.empty_like(q)
@@ -329,8 +325,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,
scale=1.0 / (head_dim**0.5),
has_text=has_text,
text_q=False,
# BLOCK_Q=BLOCK_Q,
# BLOCK_KV=BLOCK_KV,
BLOCK_DIM=BLOCK_DIM,
)
@@ -363,8 +357,6 @@ def sliding_tile_attention_triton(q: torch.Tensor,
scale=1.0 / (head_dim**0.5),
has_text=has_text,
text_q=True,
# BLOCK_Q=BLOCK_Q,
# BLOCK_KV=BLOCK_KV,
BLOCK_DIM=BLOCK_DIM,
)
fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py
-44
View File
@@ -129,21 +129,17 @@ def _rms_norm_fwd_fused(
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
# Compute variance
_var = x * x
var = tl.sum(_var, axis=1) / N
rstd = 1 / tl.sqrt(var + eps)
# Write mean / rstd
tl.store(Rstd + rows, rstd)
rstd = tl.reshape(rstd, (BLOCK_M, 1))
# Normalize and apply linear transformation
w = tl.load(W + cols)
x_hat = x * rstd
y = x_hat * w
# Write output
y = y.to(Y.type.element_ty)
tl.store(y_ptr, y, mask=mask[None, :])
@@ -181,7 +177,6 @@ def rmsnorm(x, w, eps):
else:
BLOCK_M = 1
# Call the triton kernel
_rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( #
x,
y,
@@ -229,20 +224,16 @@ def _layer_norm_param_fwd_fused(
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
# Compute mean and Variance
mean = tl.sum(x, axis=1, keep_dims=True) / N
# Compute variance
_var = (x - mean) * (x - mean)
var = tl.sum(_var, axis=1, keep_dims=True) / N
rstd = 1 / tl.sqrt(var + eps)
# Write mean / rstd
_mean = tl.reshape(mean, (BLOCK_M))
_rstd = tl.reshape(rstd, (BLOCK_M))
tl.store(Mean + rows, _mean)
tl.store(Rstd + rows, _rstd)
# Normalize and apply linear transformation
x_hat = (x - mean) * rstd
w = tl.load(W + cols)
@@ -250,7 +241,6 @@ def _layer_norm_param_fwd_fused(
x_hat = x_hat * w + b
# Write output
x_hat = x_hat.to(Y.type.element_ty)
tl.store(y_ptr, x_hat, mask=mask[None, :])
@@ -298,9 +288,6 @@ def layernorm_param(x, w, b, eps):
return y, mean, rstd
########################################################
# Elementwise_affine=False
########################################################
@triton.jit
@@ -327,23 +314,18 @@ def _layer_norm_noparam_fwd_fused(
x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32)
# Compute mean and Variance
mean = tl.sum(x, axis=1, keep_dims=True) / N
# Compute variance
_var = (x - mean) * (x - mean)
var = tl.sum(_var, axis=1, keep_dims=True) / N
rstd = 1 / tl.sqrt(var + eps)
# Write mean / rstd
_mean = tl.reshape(mean, (BLOCK_M))
_rstd = tl.reshape(rstd, (BLOCK_M))
tl.store(Mean + rows, _mean)
tl.store(Rstd + rows, _rstd)
# Normalize and apply linear transformation
x_hat = (x - mean) * rstd
# Write output
x_hat = x_hat.to(Y.type.element_ty)
tl.store(y_ptr, x_hat, mask=mask[None, :])
@@ -440,7 +422,6 @@ class Int8LinearFunction(torch.autograd.Function):
grad_bias = None
if ctx.needs_input_grad[0]:
# grad_input = grad_output @ W
w_float = dequantize_weight(w_q, w_s, grad_output.dtype)
grad_input = torch.matmul(grad_output, w_float)
@@ -489,34 +470,15 @@ class RMSNormFunction(torch.autograd.Function):
grad_output_flat = grad_output
dx = dx
# Standard RMSNorm backward
# dy = grad_output_flat
# x_hat = x * rstd
# w * dy
# c1 = mean(dy * w * x) * rstd^2
# dx = (dy * w - c1 * x) * rstd
# More precise:
# y = x * rstd * w
# dy = grad_output
# dw = sum(dy * x * rstd, dim=0)
grad_w = None
if ctx.needs_input_grad[1]: # w
# grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0)
# More accurate gradient for weight when considering rstd was computed from x
# Actually, for Affine part (w), it is just sum(dL/dy * x_hat).
# y = x_hat * w
# dL/dw = sum(dL/dy * x_hat)
x_hat = x_flat * rstd
grad_w = (grad_output_flat * x_hat).sum(dim=0)
grad_input = None
if ctx.needs_input_grad[0]: # x
# x_hat = x * rstd
# dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat)
# but for RMSNorm:
# dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula
# Using PyTorch autograd for reference logic:
# y = x / sqrt(mean(x^2) + eps) * w
@@ -530,8 +492,6 @@ class RMSNormFunction(torch.autograd.Function):
# Let's compute:
dy_w = grad_output_flat * w
# term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x)
# grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2)
# Re-derivation:
# x_hat = x * rstd
@@ -586,10 +546,6 @@ class LayerNormFunction(torch.autograd.Function):
else:
w_eff = w
# dx calculation
# x_hat = (x - mean) * rstd
# y = x_hat * w + b
# dL/dx_hat = dL/dy * w
dy = grad_output_flat
dx_hat = dy * w_eff
fastvideo-kernel/python/fastvideo_kernel/vmoba.py
+5 -15
View File
@@ -5,7 +5,7 @@ import random
import time
import torch
try:
from flash_attn import flash_attn_varlen_func # Use the new flash attention function
from flash_attn import flash_attn_varlen_func
from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
except ImportError:
@@ -46,7 +46,6 @@ def calc_chunks(cu_seqlen, moba_chunk_size):
chunk_to_batch[cu_num_chunk[1:-1]] = 1
chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)
# Do not filter out any chunk
filtered_chunk_indices = torch.arange(num_chunk, device=cu_seqlen.device, dtype=torch.int32)
num_filtered_chunk = num_chunk
@@ -66,7 +65,7 @@ def _select_threshold_query_head(gate: torch.Tensor, valid_gate_mask: torch.Tens
eps = 1e-6
# LSE‐style normalization per <head, query> (across chunks)
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
row_min = gate_min_val.amin(dim=0) # (H, S)
@@ -127,7 +126,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, g
eps = 1e-6
# LSE‐style normalization per block (across heads and queries)
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1)
@@ -193,7 +192,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
eps = 1e-6
# LSE‐style normalization globally across all valid entries
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
overall_max = gate_masked.max() # scalar
@@ -251,7 +250,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor,
# 8) create the mask in the original flat shape
others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,)
if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected
if selected_other_indices.numel() > 0:
others_mask_flat[selected_other_indices] = True
others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
@@ -349,8 +348,6 @@ class MixedAttention(torch.autograd.Function):
ctx.moba_chunk_size = moba_chunk_size
ctx.softmax_scale = softmax_scale = q.shape[-1]**(-0.5)
# Non-causal self-attention branch
# return out, softmax_lse, S_dmask, rng_state
self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
q=q,
k=k,
@@ -532,12 +529,10 @@ def moba_attn_varlen(
kv = torch.stack((k, v), dim=1)
seqlen, num_head, head_dim = q.shape
# Compute chunk boundaries.
cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(cu_seqlens, moba_chunk_size)
self_attn_cu_seqlen = cu_chunk
# Update top-k selection to include the self chunk.
moba_topk = min(moba_topk, num_filtered_chunk)
# --- Build filtered KV from chunks ---
@@ -553,16 +548,13 @@ def moba_attn_varlen(
gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)
# Compute key_gate_weight over valid tokens.
key_values = gathered[:, :, 0].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
key_sum = (key_values * valid_mask_exp).sum(dim=1)
divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
key_gate_weight = key_sum / divisor # [num_filtered_chunk, num_head, head_dim]
# Compute gate logits between key_gate_weight and queries.
q_float = q.float()
# gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen]
gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
# Amplify the diagonal (self chunk) contributions.
@@ -650,7 +642,6 @@ def moba_attn_varlen(
moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
# Build cumulative sequence lengths for the selected queries.
moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
q_zero_mask = moba_seqlen_q == 0
valid_expert_mask = ~q_zero_mask
@@ -814,7 +805,6 @@ def test_attn_varlen_moba_speed(batch,
simsum_threshold=0.25,
threshold_type='query_head'):
"""Speed test comparing flash_attn vs moba_attention"""
# Get data
q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
print(
f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}"
fastvideo-kernel/tests/test_attn_qat_train.py
-143
View File
@@ -53,15 +53,12 @@ def attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=Non
if sm_scale is None:
sm_scale = 1.0 / sqrt(q_BLHD.shape[-1])
# Convert from BLHD to BHLD format
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
# Call attention with BHLD format
o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale)
# Convert back from BHLD to BLHD format
return o_BHLD.permute(0, 2, 1, 3).contiguous()
@@ -84,7 +81,6 @@ def fused_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None, w
if sm_scale is None:
sm_scale = 1.0 / sqrt(q_BLHD.shape[-1])
# Convert from BLHD to BHLD format
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -93,7 +89,6 @@ def fused_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None, w
# Note: fused_attention expects inputs in BHLD format and only supports same sequence lengths
o_BHLD = fused_attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale, warp_specialize)
# Convert back from BHLD to BLHD format
return o_BHLD.permute(0, 2, 1, 3).contiguous()
@@ -112,7 +107,6 @@ def cosine_similarity(tensor1, tensor2):
t1_flat = tensor1.flatten().float()
t2_flat = tensor2.flatten().float()
# Compute cosine similarity: (A · B) / (||A|| * ||B||)
dot_product = torch.dot(t1_flat, t2_flat)
norm1 = torch.norm(t1_flat)
norm2 = torch.norm(t2_flat)
@@ -139,7 +133,6 @@ def naive_attention(q, k, v, causal, sm_scale):
Returns:
Output tensor of shape (Z, H, N_CTX_Q, HEAD_DIM)
"""
# Compute attention scores: QK^T
scores = torch.matmul(q, k.transpose(-2, -1)) * sm_scale
# Apply causal mask if needed
@@ -156,10 +149,8 @@ def naive_attention(q, k, v, causal, sm_scale):
# In cross-attention, typically no causal masking is applied
pass
# Apply softmax
attn_weights = torch.softmax(scores, dim=-1)
# Apply attention weights to values
out = torch.matmul(attn_weights, v)
return out
@@ -175,7 +166,6 @@ def test_qat_attention_forward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = False
# Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
@@ -183,18 +173,15 @@ def test_qat_attention_forward():
# Test with QAT (using wrapper that does permute/contiguous)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention comparison
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale)
out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Check that outputs have correct shape
assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
# Check that outputs are finite
assert torch.isfinite(out_qat_BLHD).all()
assert torch.isfinite(out_naive_BLHD).all()
@@ -206,9 +193,6 @@ def test_qat_attention_forward():
f" QAT vs Naive Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# QAT output should be reasonably close to naive (within tolerance due to quantization)
# Using a reasonable tolerance for float16 and quantization effects
# assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}"
print("✓ Forward pass test passed.")
@@ -223,12 +207,10 @@ def test_qat_attention_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for QAT
q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
# Create input tensors for naive (same values, convert to BHLD)
q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
@@ -250,7 +232,6 @@ def test_qat_attention_backward():
# Backward pass for naive
out_naive_BLHD.backward(dout)
# Check that gradients are computed
assert q_qat_BLHD.grad is not None
assert k_qat_BLHD.grad is not None
assert v_qat_BLHD.grad is not None
@@ -258,17 +239,14 @@ def test_qat_attention_backward():
assert k_naive_BHLD.grad is not None
assert v_naive_BHLD.grad is not None
# Check that gradients have correct shape
assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape
assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape
assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape
# Check that gradients are finite
assert torch.isfinite(q_qat_BLHD.grad).all()
assert torch.isfinite(k_qat_BLHD.grad).all()
assert torch.isfinite(v_qat_BLHD.grad).all()
# Convert naive gradients to BLHD format for comparison
q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
@@ -292,10 +270,6 @@ def test_qat_attention_backward():
f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}"
)
# Gradients should be reasonably close (within tolerance due to quantization)
# assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}"
# assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}"
# assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}"
print("✓ Backward pass test passed.")
@@ -313,7 +287,6 @@ def test_qat_attention_different_shapes():
causal = True
for Z, H, N_CTX, HEAD_DIM in test_configs:
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
@@ -321,7 +294,6 @@ def test_qat_attention_different_shapes():
sm_scale = 1.0 / sqrt(HEAD_DIM)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -341,8 +313,6 @@ def test_qat_attention_different_shapes():
f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}"
print(f"✓ Shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})")
@@ -356,14 +326,12 @@ def test_qat_attention_non_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = False
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -383,8 +351,6 @@ def test_qat_attention_non_causal():
f" QAT vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}"
print("✓ Non-causal test passed.")
@@ -398,14 +364,12 @@ def test_qat_attention_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -425,8 +389,6 @@ def test_qat_attention_causal():
f" QAT vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}"
print("✓ Causal test passed.")
@@ -440,12 +402,10 @@ def test_qat_attention_causal_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for QAT
q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
# Create input tensors for naive (same values, convert to BHLD)
q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
@@ -457,7 +417,6 @@ def test_qat_attention_causal_backward():
out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale)
out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Create dummy gradient
dout = torch.randn_like(out_qat_BLHD).contiguous()
# Backward pass for QAT
@@ -466,12 +425,10 @@ def test_qat_attention_causal_backward():
# Backward pass for naive
out_naive_BHLD.backward(dout)
# Check that gradients are computed
assert q_qat_BLHD.grad is not None
assert k_qat_BLHD.grad is not None
assert v_qat_BLHD.grad is not None
# Convert naive gradients to BLHD format for comparison
q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
@@ -513,7 +470,6 @@ def test_qat_attention_different_seq_lengths():
for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs:
sm_scale = 1.0 / sqrt(HEAD_DIM)
# Create input tensors in BLHD format with different sequence lengths
q_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE)
@@ -521,7 +477,6 @@ def test_qat_attention_different_seq_lengths():
# Test with QAT (using wrapper that does permute/contiguous)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention comparison
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -532,7 +487,6 @@ def test_qat_attention_different_seq_lengths():
assert out_qat_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM)
assert out_naive_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM)
# Check that outputs are finite
assert torch.isfinite(out_qat_BLHD).all()
assert torch.isfinite(out_naive_BLHD).all()
@@ -544,8 +498,6 @@ def test_qat_attention_different_seq_lengths():
f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}, causal={causal}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}), got max_diff={max_diff.item():.6f}"
print(f"✓ Different sequence lengths test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})")
@@ -564,12 +516,10 @@ def test_qat_attention_different_seq_lengths_backward():
for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs:
sm_scale = 1.0 / sqrt(HEAD_DIM)
# Create input tensors in BLHD format with different sequence lengths
q_qat_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
# Create input tensors for naive (same values, convert to BHLD)
q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
@@ -581,7 +531,6 @@ def test_qat_attention_different_seq_lengths_backward():
out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale)
out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Create dummy gradient
dout = torch.randn_like(out_qat_BLHD).contiguous()
# Backward pass for QAT
@@ -590,7 +539,6 @@ def test_qat_attention_different_seq_lengths_backward():
# Backward pass for naive
out_naive_BHLD.backward(dout)
# Check that gradients are computed
assert q_qat_BLHD.grad is not None
assert k_qat_BLHD.grad is not None
assert v_qat_BLHD.grad is not None
@@ -598,17 +546,14 @@ def test_qat_attention_different_seq_lengths_backward():
assert k_naive_BHLD.grad is not None
assert v_naive_BHLD.grad is not None
# Check that gradients have correct shape
assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape
assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape
assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape
# Check that gradients are finite
assert torch.isfinite(q_qat_BLHD.grad).all()
assert torch.isfinite(k_qat_BLHD.grad).all()
assert torch.isfinite(v_qat_BLHD.grad).all()
# Convert naive gradients to BLHD format for comparison
q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
@@ -632,10 +577,6 @@ def test_qat_attention_different_seq_lengths_backward():
f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dV: Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}"
)
# Gradients should be reasonably close (within tolerance due to quantization)
# assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}"
# assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}"
# assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}"
print(f"✓ Cross-attention backward test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})")
@@ -650,7 +591,6 @@ def test_fused_attention_forward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = False
# Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
@@ -658,18 +598,15 @@ def test_fused_attention_forward():
# Test with fused attention (using wrapper that does permute/contiguous)
out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention comparison
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale)
out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Check that outputs have correct shape
assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
# Check that outputs are finite
assert torch.isfinite(out_fused_BLHD).all()
assert torch.isfinite(out_naive_BLHD).all()
@@ -681,8 +618,6 @@ def test_fused_attention_forward():
f" Fused vs Naive - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Fused attention should be reasonably close to naive (within tolerance for float16)
# assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}"
print("✓ Fused attention forward pass test passed.")
@@ -697,12 +632,10 @@ def test_fused_attention_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for fused attention
q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
# Create input tensors for naive (same values, convert to BHLD)
q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
@@ -714,7 +647,6 @@ def test_fused_attention_backward():
out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale)
out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Create dummy gradient (same for both, in BLHD format)
dout = torch.randn_like(out_fused_BLHD).contiguous()
# Backward pass for fused attention
@@ -723,7 +655,6 @@ def test_fused_attention_backward():
# Backward pass for naive
out_naive_BLHD.backward(dout)
# Check that gradients are computed
assert q_fused_BLHD.grad is not None
assert k_fused_BLHD.grad is not None
assert v_fused_BLHD.grad is not None
@@ -731,17 +662,14 @@ def test_fused_attention_backward():
assert k_naive_BHLD.grad is not None
assert v_naive_BHLD.grad is not None
# Check that gradients have correct shape
assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape
assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape
assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape
# Check that gradients are finite
assert torch.isfinite(q_fused_BLHD.grad).all()
assert torch.isfinite(k_fused_BLHD.grad).all()
assert torch.isfinite(v_fused_BLHD.grad).all()
# Convert naive gradients to BLHD format for comparison
q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
@@ -765,10 +693,6 @@ def test_fused_attention_backward():
f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}"
)
# Gradients should be reasonably close (within tolerance for float16)
# assert dq_diff.max() < 1e-1, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}"
# assert dk_diff.max() < 1e-1, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}"
# assert dv_diff.max() < 1e-1, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}"
print("✓ Fused attention backward pass test passed.")
@@ -787,7 +711,6 @@ def test_fused_attention_different_shapes():
causal = True
for Z, H, N_CTX, HEAD_DIM in test_configs:
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
@@ -795,7 +718,6 @@ def test_fused_attention_different_shapes():
sm_scale = 1.0 / sqrt(HEAD_DIM)
out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -815,8 +737,6 @@ def test_fused_attention_different_shapes():
f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}"
print(f"✓ Fused attention shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})")
@@ -830,14 +750,12 @@ def test_fused_attention_non_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = False
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -857,8 +775,6 @@ def test_fused_attention_non_causal():
f" Fused vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}"
print("✓ Fused attention non-causal test passed.")
@@ -872,14 +788,12 @@ def test_fused_attention_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -899,8 +813,6 @@ def test_fused_attention_causal():
f" Fused vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs should be reasonably close
# assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}"
print("✓ Fused attention causal test passed.")
@@ -914,12 +826,10 @@ def test_fused_attention_causal_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for fused attention
q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True)
# Create input tensors for naive (same values, convert to BHLD)
q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True)
@@ -931,7 +841,6 @@ def test_fused_attention_causal_backward():
out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale)
out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# Create dummy gradient
dout = torch.randn_like(out_fused_BLHD).contiguous()
# Backward pass for fused attention
@@ -940,12 +849,10 @@ def test_fused_attention_causal_backward():
# Backward pass for naive
out_naive_BHLD.backward(dout)
# Check that gradients are computed
assert q_fused_BLHD.grad is not None
assert k_fused_BLHD.grad is not None
assert v_fused_BLHD.grad is not None
# Convert naive gradients to BLHD format for comparison
q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous()
@@ -996,14 +903,11 @@ def test_fused_vs_qat_attention_forward():
v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16)
out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale)
# Convert QAT output back to float16 for comparison
out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16)
# Check that outputs have correct shape
assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
# Check that outputs are finite
assert torch.isfinite(out_fused_BLHD).all()
assert torch.isfinite(out_qat_BLHD).all()
@@ -1015,8 +919,6 @@ def test_fused_vs_qat_attention_forward():
f" Fused vs QAT - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# They may differ due to quantization in QAT and different implementations
# assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close, got max_diff={max_diff.item():.6f}"
print("✓ Fused vs QAT forward pass test passed.")
@@ -1030,12 +932,10 @@ def test_fused_vs_qat_attention_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for fused attention (float16)
q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
# Create input tensors for QAT (float16, same values)
q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
@@ -1046,7 +946,6 @@ def test_fused_vs_qat_attention_backward():
# Forward pass with QAT
out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale)
# Create dummy gradient (same for both, in BLHD format)
dout = torch.randn_like(out_fused_BLHD).contiguous()
dout_qat = dout.to(torch.bfloat16)
@@ -1056,7 +955,6 @@ def test_fused_vs_qat_attention_backward():
# Backward pass for QAT
out_qat_BLHD.backward(dout_qat)
# Check that gradients are computed
assert q_fused_BLHD.grad is not None
assert k_fused_BLHD.grad is not None
assert v_fused_BLHD.grad is not None
@@ -1064,17 +962,14 @@ def test_fused_vs_qat_attention_backward():
assert k_qat_BLHD.grad is not None
assert v_qat_BLHD.grad is not None
# Check that gradients have correct shape
assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape
assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape
assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape
# Check that gradients are finite
assert torch.isfinite(q_fused_BLHD.grad).all()
assert torch.isfinite(k_fused_BLHD.grad).all()
assert torch.isfinite(v_fused_BLHD.grad).all()
# Convert QAT gradients to float16 for comparison
q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16)
k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16)
v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16)
@@ -1098,10 +993,6 @@ def test_fused_vs_qat_attention_backward():
f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}"
)
# Gradients may differ due to quantization in QAT and different implementations
# assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}"
# assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}"
# assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}"
print("✓ Fused vs QAT backward pass test passed.")
@@ -1120,7 +1011,6 @@ def test_fused_vs_qat_attention_different_shapes():
for Z, H, N_CTX, HEAD_DIM in test_configs:
sm_scale = 1.0 / sqrt(HEAD_DIM)
# Create input tensors in BLHD format (float16)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
@@ -1148,8 +1038,6 @@ def test_fused_vs_qat_attention_different_shapes():
f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs may differ due to quantization in QAT
# assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}"
print(f"✓ Fused vs QAT shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})")
@@ -1162,7 +1050,6 @@ def test_fused_vs_qat_attention_non_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = False
# Create input tensors in BLHD format (float16)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
@@ -1190,8 +1077,6 @@ def test_fused_vs_qat_attention_non_causal():
f" Fused vs QAT (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs may differ due to quantization in QAT
# assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}"
print("✓ Fused vs QAT non-causal test passed.")
@@ -1204,7 +1089,6 @@ def test_fused_vs_qat_attention_causal():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format (float16)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE)
@@ -1232,8 +1116,6 @@ def test_fused_vs_qat_attention_causal():
f" Fused vs QAT (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}"
)
# Outputs may differ due to quantization in QAT
# assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}"
print("✓ Fused vs QAT causal test passed.")
@@ -1246,12 +1128,10 @@ def test_fused_vs_qat_attention_causal_backward():
sm_scale = 1.0 / sqrt(HEAD_DIM)
causal = True
# Create input tensors in BLHD format for fused attention (float16)
q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True)
# Create input tensors for QAT (float16, same values)
q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True)
@@ -1262,7 +1142,6 @@ def test_fused_vs_qat_attention_causal_backward():
# Forward pass with QAT
out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale)
# Create dummy gradient
dout = torch.randn_like(out_fused_BLHD).contiguous()
dout_qat = dout.to(torch.bfloat16)
@@ -1272,7 +1151,6 @@ def test_fused_vs_qat_attention_causal_backward():
# Backward pass for QAT
out_qat_BLHD.backward(dout_qat)
# Check that gradients are computed
assert q_fused_BLHD.grad is not None
assert k_fused_BLHD.grad is not None
assert v_fused_BLHD.grad is not None
@@ -1280,7 +1158,6 @@ def test_fused_vs_qat_attention_causal_backward():
assert k_qat_BLHD.grad is not None
assert v_qat_BLHD.grad is not None
# Convert QAT gradients to float16 for comparison
q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16)
k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16)
v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16)
@@ -1319,7 +1196,6 @@ def test_qat_attention_wan_shape_forward():
print(f" Testing WAN shape: Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}")
# Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM)
q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE)
@@ -1327,7 +1203,6 @@ def test_qat_attention_wan_shape_forward():
# Test with QAT (using wrapper that does permute/contiguous)
out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale)
# Convert to BHLD for naive attention comparison
q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous()
k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous()
v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
@@ -1338,11 +1213,9 @@ def test_qat_attention_wan_shape_forward():
print(f" Expected shape: {(Z, N_CTX, H, HEAD_DIM)}")
print(f" out_qat_BLHD shape: {out_qat_BLHD.shape}")
# Check that outputs have correct shape
assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM)
# Check that outputs are finite
assert torch.isfinite(out_qat_BLHD).all()
assert torch.isfinite(out_naive_BLHD).all()
@@ -1393,7 +1266,6 @@ def test_qat_attention_wan_shape_backward():
# -----------------------------
out_naive_BHLD = naive_attention(q_BHLD, k_BHLD, v_BHLD, causal, sm_scale)
# Convert naive output back to BLHD for comparing outputs
out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous()
# -----------------------------
@@ -1597,23 +1469,8 @@ if __name__ == "__main__":
print("Running Fused attention tests...")
print()
# test_fused_attention_forward()
# test_fused_attention_backward()
# test_fused_attention_different_shapes()
# test_fused_attention_non_causal()
# test_fused_attention_causal()
# test_fused_attention_causal_backward()
# print()
# print("Running Fused vs QAT attention comparison tests...")
# print()
# test_fused_vs_qat_attention_forward()
# test_fused_vs_qat_attention_backward()
# test_fused_vs_qat_attention_different_shapes()
# test_fused_vs_qat_attention_non_causal()
# test_fused_vs_qat_attention_causal()
# test_fused_vs_qat_attention_causal_backward()
print()
print("All tests passed! ✓")
fastvideo-kernel/tests/test_sta.py
-2
View File
@@ -2,7 +2,6 @@ import torch
from .support_flex_sta import get_sliding_tile_attention_mask
from fastvideo_kernel import sliding_tile_attention
from torch.nn.attention.flex_attention import flex_attention
# from flash_attn_interface import flash_attn_func
from tqdm import tqdm
flex_attention = torch.compile(flex_attention, dynamic=False)
@@ -83,7 +82,6 @@ def test_sliding_tile_attention():
mean = 1e-1
std = 10
# Run correctness check directly
results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
assert results['TK vs FLEX'][
'avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
fastvideo-kernel/tests/test_turbodiffusion.py
-12
View File
@@ -3,8 +3,6 @@ import pytest
import sys
import os
# Ensure local package is imported
# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../python")))
try:
from fastvideo_kernel._C import fastvideo_kernel_ops
@@ -65,20 +63,13 @@ class TestTurboDiffusion:
M, N, K = 32, 64, 128
x = torch.randn(M, K, dtype=dtype, device="cuda")
# Create weights
# For simplicity in testing, let's create random int8 weights and scales
w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda")
# Scale shape: The Int8Linear class uses:
# row_blocks = cdiv(out_features, b=128)
# col_blocks = cdiv(in_features, b=128)
# scale shape: (row_blocks, col_blocks)
row_blocks = (N + 127) // 128
col_blocks = (K + 127) // 128
w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs()
# Run int8_linear
output = turbodiffusion_ops.int8_linear(x, w_q, w_s)
assert output.shape == (M, N)
@@ -92,7 +83,6 @@ class TestTurboDiffusion:
M, N, K = 32, 64, 128
x = torch.randn(M, K, dtype=dtype, device="cuda", requires_grad=True)
# Weights (frozen)
w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda")
row_blocks = (N + 127) // 128
col_blocks = (K + 127) // 128
@@ -100,7 +90,6 @@ class TestTurboDiffusion:
bias = torch.randn(N, dtype=dtype, device="cuda", requires_grad=True)
# Use Int8LinearFunction
output = turbodiffusion_ops.Int8LinearFunction.apply(x, w_q, w_s, bias)
loss = output.sum()
loss.backward()
@@ -109,7 +98,6 @@ class TestTurboDiffusion:
assert bias.grad is not None
assert x.grad.shape == (M, K)
# Check correctness against dequantized weight
w_float = turbodiffusion_ops.dequantize_weight(w_q, w_s, dtype)
# Reference

Some files were not shown because too many files have changed in this diff Show More