From 65058496eec70d21ae8add17caaee589e42e683d Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Sun, 19 Oct 2025 08:36:52 -0400 Subject: [PATCH 1/4] fix: diffusion bottleneck --- scGraphLLM/graph_op.py | 52 +++++++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 18 deletions(-) diff --git a/scGraphLLM/graph_op.py b/scGraphLLM/graph_op.py index e28c331..ef6c478 100644 --- a/scGraphLLM/graph_op.py +++ b/scGraphLLM/graph_op.py @@ -1,8 +1,9 @@ import torch from torch_geometric.utils import scatter, remove_self_loops - +from torch_geometric.data import Data, Batch from scGraphLLM._globals import * ## imported global variables are all caps + def _identity(x): return x @@ -58,7 +59,7 @@ def _chebyshev_coeff(L_rescaled, K, func, N=100): return c_k @torch.amp.autocast(enabled=False, device_type='cuda') -def _chebyshev_diffusion_per_sample(edge_index, num_nodes, E, k=128, edge_weight=None, beta=0.5): +def _chebyshev_diffusion_batch(edge_index, num_nodes, E, k=128, edge_weight=None, beta=0.5): """ E: (S, H, d) """ @@ -94,20 +95,35 @@ def _chebyshev_diffusion(edge_index_list, num_nodes_list, E, k=64, beta=0.5): E: (B, S, H, d) """ B, S, H, D = E.size() - final_emb = [] - for i in range(B): - E_i = E[i, :num_nodes_list[i], ...] - edge_index = edge_index_list[i] - num_nodes = num_nodes_list[i] - sample_emb = _chebyshev_diffusion_per_sample(edge_index, num_nodes_list[i], E_i, k=k, beta=beta) - - # pad zero at the right end - pad_size = S - sample_emb.size(0) - if pad_size > 0: - zero_pad_right = torch.zeros(pad_size, H, D, device=E.device, dtype=E.dtype) - sample_emb = torch.cat([sample_emb, zero_pad_right], dim=0) - final_emb.append(sample_emb) - fe = torch.stack(final_emb, dim=0) - assert fe.size() == E.size(), f"Expect {E.size()}, Got {fe.size()}" - return fe \ No newline at end of file + # 1. Create a mask to un-pad the (B, S, H, D) tensor + num_nodes_tensor = torch.tensor(num_nodes_list, device=E.device, dtype=torch.long) + mask = torch.arange(S, device=E.device)[None, :] < num_nodes_tensor[:, None] + + # 2. Un-pad E + # Shape goes from (B, S, H, D) -> (total_nodes, H, D) + unpadded_E = E[mask] + + # 3. Create a single PyG Batch object + data_list = [Data(edge_index=ei, num_nodes=n) for ei, n in zip(edge_index_list, num_nodes_list)] + pyg_batch = Batch.from_data_list(data_list) + + batched_edge_index = pyg_batch.edge_index.to(E.device) + total_nodes = pyg_batch.num_nodes # This is sum(num_nodes_list) + + # 4. Run batched diffusion + diffused_unpadded_E = _chebyshev_diffusion_batch( + batched_edge_index, + total_nodes, + unpadded_E, + k=k, + edge_weight=None, # Assuming edge_weight is None + beta=beta + ) + + # 5. Re-pad the result + final_emb = torch.zeros_like(E) + final_emb[mask] = diffused_unpadded_E + + assert final_emb.size() == E.size(), f"Expect {E.size()}, Got {final_emb.size()}" + return final_emb \ No newline at end of file From f14850cdf38b4887b86f20e6984b5928c0efb77c Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Sat, 1 Nov 2025 18:43:59 -0400 Subject: [PATCH 2/4] fix: inference bottlenecks --- scGraphLLM/data.py | 3 ++- scGraphLLM/inference.py | 41 +++++++++++++++++++++++++++++------------ 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/scGraphLLM/data.py b/scGraphLLM/data.py index 5c49e3f..c167187 100755 --- a/scGraphLLM/data.py +++ b/scGraphLLM/data.py @@ -58,8 +58,9 @@ def scglm_collate_fn(batch, pad_node, inference=False): data["obs_name"] = [] # Make a dictionary of lists from the list of dictionaries + keys = list(data.keys()) for b in batch: - for key in data.keys(): + for key in keys: data[key].append(b[key]) # Pad these dictionaries of lists diff --git a/scGraphLLM/inference.py b/scGraphLLM/inference.py index d3af2f3..5ad74a4 100644 --- a/scGraphLLM/inference.py +++ b/scGraphLLM/inference.py @@ -168,9 +168,10 @@ def get_cell_embeddings( x_list = [] obs_names_list = [] + device = next(model.parameters()).device with torch.no_grad(): for batch in tqdm(dataloader, desc="Forward Pass"): - seq_lengths = torch.tensor(batch["num_nodes"]).to("cuda") + seq_lengths = torch.tensor(batch["num_nodes"], device=device) obs_names = batch["obs_name"] x = model(send_to_gpu(batch))[0] # shape: [B, T, H] gene_ids = batch["orig_gene_id"] # shape: [B, T] @@ -234,19 +235,35 @@ def get_gene_embeddings( gene_counts = defaultdict(int) with torch.no_grad(): for batch in tqdm(dataloader, desc="Forward Pass"): - seq_lengths = torch.tensor(batch["num_nodes"]).to("cuda") - gene_ids = batch["orig_gene_id"].detach().cpu().numpy() + seq_lengths = torch.tensor(batch["num_nodes"], device="cuda") + gene_ids = batch["orig_gene_id"].to("cuda") # Keep on GPU initially x = model(send_to_gpu(batch))[0] # shape: [B, T, H] - for x_cell, ids, seq_len in zip(x, gene_ids, seq_lengths): - for j in range(seq_len): - gene_id = ids[j] - x_gene = x_cell[j].detach().cpu().numpy() - if gene_id not in gene_embedding_sums: - gene_embedding_sums[gene_id] = x_gene.copy() - else: - gene_embedding_sums[gene_id] += x_gene - gene_counts[gene_id] += 1 + # Vectorized approach: process all valid tokens at once + batch_size, max_seq_len = x.shape[:2] + # Create a mask for valid positions + valid_mask = torch.arange(max_seq_len, device=x.device)[None, :] < seq_lengths[:, None] # [B, T] + + # Flatten and extract valid embeddings + flat_x = x.reshape(-1, x.shape[-1]) # [B*T, H] + flat_gene_ids = gene_ids.reshape(-1) # [B*T] + flat_mask = valid_mask.reshape(-1) # [B*T] + + # Get valid embeddings only + valid_x = flat_x[flat_mask] # [num_valid, H] + valid_gene_ids = flat_gene_ids[flat_mask].cpu().numpy() # [num_valid] + + # Process in chunks to avoid memory issues for very large batches + # Group by gene_id and accumulate embeddings + for gene_id in np.unique(valid_gene_ids): + mask = valid_gene_ids == gene_id + gene_emb = valid_x[mask].sum(dim=0).detach().cpu().numpy() + + if gene_id not in gene_embedding_sums: + gene_embedding_sums[gene_id] = gene_emb + else: + gene_embedding_sums[gene_id] += gene_emb + gene_counts[gene_id] += mask.sum().item() # compute average embedding per gene gene_embeddings = { From 2e745bfc68bb233ec8e07a85ce474d2035a0857b Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Sat, 1 Nov 2025 21:17:16 -0400 Subject: [PATCH 3/4] much better inference step --- scGraphLLM/inference.py | 63 +++++++++++++++++++++++------------------ 1 file changed, 35 insertions(+), 28 deletions(-) diff --git a/scGraphLLM/inference.py b/scGraphLLM/inference.py index 5ad74a4..276a105 100644 --- a/scGraphLLM/inference.py +++ b/scGraphLLM/inference.py @@ -235,35 +235,42 @@ def get_gene_embeddings( gene_counts = defaultdict(int) with torch.no_grad(): for batch in tqdm(dataloader, desc="Forward Pass"): - seq_lengths = torch.tensor(batch["num_nodes"], device="cuda") - gene_ids = batch["orig_gene_id"].to("cuda") # Keep on GPU initially - x = model(send_to_gpu(batch))[0] # shape: [B, T, H] - - # Vectorized approach: process all valid tokens at once - batch_size, max_seq_len = x.shape[:2] - # Create a mask for valid positions - valid_mask = torch.arange(max_seq_len, device=x.device)[None, :] < seq_lengths[:, None] # [B, T] - - # Flatten and extract valid embeddings - flat_x = x.reshape(-1, x.shape[-1]) # [B*T, H] - flat_gene_ids = gene_ids.reshape(-1) # [B*T] - flat_mask = valid_mask.reshape(-1) # [B*T] - - # Get valid embeddings only - valid_x = flat_x[flat_mask] # [num_valid, H] - valid_gene_ids = flat_gene_ids[flat_mask].cpu().numpy() # [num_valid] - - # Process in chunks to avoid memory issues for very large batches - # Group by gene_id and accumulate embeddings - for gene_id in np.unique(valid_gene_ids): - mask = valid_gene_ids == gene_id - gene_emb = valid_x[mask].sum(dim=0).detach().cpu().numpy() - - if gene_id not in gene_embedding_sums: - gene_embedding_sums[gene_id] = gene_emb + seq_lengths = torch.as_tensor(batch["num_nodes"], device="cuda", dtype=torch.long) + gene_ids = batch["orig_gene_id"].to("cuda") # [B, T] + x = model(send_to_gpu(batch))[0] # [B, T, H] + + B, T, H = x.shape + valid_mask = torch.arange(T, device=x.device)[None, :] < seq_lengths[:, None] # [B, T] + + flat_x = x.reshape(-1, H) # [B*T, H] + flat_gene_ids = gene_ids.reshape(-1).to(torch.long) # [B*T] + flat_mask = valid_mask.reshape(-1) # [B*T] + + # Keep only valid tokens + valid_x = flat_x[flat_mask] # [N, H] + valid_gene_ids = flat_gene_ids[flat_mask] # [N] + + # Group-by gene_id on GPU + unique_ids, inverse = torch.unique(valid_gene_ids, return_inverse=True) + + # Sum embeddings per unique gene id + sums = torch.zeros((unique_ids.numel(), H), device=valid_x.device, dtype=valid_x.dtype) + sums.index_add_(0, inverse, valid_x) # segment sum + + # Counts per unique id + counts = torch.bincount(inverse, minlength=unique_ids.numel()) # [U] + + # Move once to CPU and update Python dicts + unique_ids_cpu = unique_ids.tolist() + sums_cpu = sums.cpu().numpy() + counts_cpu = counts.cpu().tolist() + + for gid, emb_sum, cnt in zip(unique_ids_cpu, sums_cpu, counts_cpu): + if gid not in gene_embedding_sums: + gene_embedding_sums[gid] = emb_sum.copy() else: - gene_embedding_sums[gene_id] += gene_emb - gene_counts[gene_id] += mask.sum().item() + gene_embedding_sums[gid] += emb_sum + gene_counts[gid] += int(cnt) # compute average embedding per gene gene_embeddings = { From 2dd97571a472ee0a0ca64b6c75e4e108db7205e2 Mon Sep 17 00:00:00 2001 From: sputti-czi Date: Thu, 18 Dec 2025 17:56:32 +0530 Subject: [PATCH 4/4] coo to csr conversion and caching coef --- scGraphLLM/graph_op.py | 106 ++++++++++++++++++++++++++++------------- 1 file changed, 74 insertions(+), 32 deletions(-) diff --git a/scGraphLLM/graph_op.py b/scGraphLLM/graph_op.py index ef6c478..2d9d70f 100644 --- a/scGraphLLM/graph_op.py +++ b/scGraphLLM/graph_op.py @@ -47,46 +47,73 @@ def _rescaled_L(edge_index, num_nodes, edge_weight=None): L_rescaled = torch.sparse_coo_tensor(edge_index, -edge_weight, (num_nodes, num_nodes)) return L_rescaled -def _chebyshev_coeff(L_rescaled, K, func, N=100): - # Gauss-Chebyshev quadrature - ind = torch.arange(0, K+1, dtype=torch.float32, device=L_rescaled.device) - ratio = torch.pi * (torch.arange(1, N+1, dtype=torch.float32, device=L_rescaled.device) - 0.5) / N - x = torch.cos(ratio) # quadrature points - T_kx = torch.cos(ind.view(-1, 1) * ratio) - w = torch.ones(N, device=L_rescaled.device) * (torch.pi / N) - f_x = func(x) - c_k = (2 / torch.pi) * torch.matmul(T_kx, w * f_x) - return c_k +# Cache for Chebyshev coefficients - only depends on K, beta, and device +_chebyshev_coeff_cache = {} -@torch.amp.autocast(enabled=False, device_type='cuda') -def _chebyshev_diffusion_batch(edge_index, num_nodes, E, k=128, edge_weight=None, beta=0.5): +def _get_cached_chebyshev_coeff(K, beta, device, N=100): """ - E: (S, H, d) + Get cached Chebyshev coefficients. These only depend on K and beta, + so we can cache them to avoid recomputation every forward pass. + """ + cache_key = (K, beta, str(device)) + + if cache_key not in _chebyshev_coeff_cache: + ind = torch.arange(0, K+1, dtype=torch.float32, device=device) + ratio = torch.pi * (torch.arange(1, N+1, dtype=torch.float32, device=device) - 0.5) / N + x = torch.cos(ratio) # quadrature points + T_kx = torch.cos(ind.view(-1, 1) * ratio) + w = torch.ones(N, device=device) * (torch.pi / N) + f_x = _exp_kernel(x, beta) + c_k = (2 / torch.pi) * torch.matmul(T_kx, w * f_x) + _chebyshev_coeff_cache[cache_key] = c_k + + return _chebyshev_coeff_cache[cache_key] + +def _chebyshev_recurrence(L_csr, E_reshaped, c_k, K): + """ + Core Chebyshev recurrence loop - separated for torch.compile optimization. """ - L_rescaled = _rescaled_L(edge_index, num_nodes, edge_weight) - c_k = _chebyshev_coeff(L_rescaled, k, lambda x: _exp_kernel(x, beta)) - E = E.to(torch.float32) - s, h, d = E.size() - assert s == num_nodes, f"Expect {num_nodes} nodes, Got {s}" - E_reshaped = E.reshape(num_nodes, h * d) - c_k = c_k.to(torch.float32) T_0 = E_reshaped - T_1 = torch.sparse.mm(L_rescaled, E_reshaped) + T_1 = torch.sparse.mm(L_csr, E_reshaped) y = c_k[0] * T_0 + c_k[1] * T_1 # start recursion T_k_prev = T_1 T_k_prev_prev = T_0 - for k in range(2, k + 1): - T_k = 2 * torch.sparse.mm(L_rescaled, T_k_prev) - T_k_prev_prev - y += c_k[k] * T_k + for i in range(2, K + 1): + T_k = 2 * torch.sparse.mm(L_csr, T_k_prev) - T_k_prev_prev + y = y + c_k[i] * T_k # shift index T_k_prev_prev = T_k_prev T_k_prev = T_k - + + return y + +@torch.amp.autocast(enabled=False, device_type='cuda') +def _chebyshev_diffusion_batch(edge_index, num_nodes, E, k=128, edge_weight=None, beta=0.5): + """ + E: (S, H, d) + """ + L_rescaled = _rescaled_L(edge_index, num_nodes, edge_weight) + + # Convert to CSR format for faster sparse matrix multiplication + L_csr = L_rescaled.to_sparse_csr() + + # Use cached coefficients (avoids recomputation every forward pass) + c_k = _get_cached_chebyshev_coeff(k, beta, E.device) + + E = E.to(torch.float32) + s, h, d = E.size() + assert s == num_nodes, f"Expect {num_nodes} nodes, Got {s}" + E_reshaped = E.reshape(num_nodes, h * d) + c_k = c_k.to(torch.float32) + + y = _chebyshev_recurrence(L_csr, E_reshaped, c_k, k) + final_emb = y.reshape(num_nodes, h, d) final_emb = final_emb.bfloat16() + return final_emb def _chebyshev_diffusion(edge_index_list, num_nodes_list, E, k=64, beta=0.5): @@ -98,32 +125,47 @@ def _chebyshev_diffusion(edge_index_list, num_nodes_list, E, k=64, beta=0.5): # 1. Create a mask to un-pad the (B, S, H, D) tensor num_nodes_tensor = torch.tensor(num_nodes_list, device=E.device, dtype=torch.long) + + # Create a (B, S) boolean mask + # Broadcasting creates a (B, S) matrix of comparisons mask = torch.arange(S, device=E.device)[None, :] < num_nodes_tensor[:, None] # 2. Un-pad E + # Use the boolean mask to select only the "real" node embeddings # Shape goes from (B, S, H, D) -> (total_nodes, H, D) unpadded_E = E[mask] - # 3. Create a single PyG Batch object - data_list = [Data(edge_index=ei, num_nodes=n) for ei, n in zip(edge_index_list, num_nodes_list)] - pyg_batch = Batch.from_data_list(data_list) + # 3. Batch edge indices with offsets (avoids PyG Batch.from_data_list overhead) + offsets = torch.zeros(B, dtype=torch.long, device=E.device) + if B > 1: + offsets[1:] = num_nodes_tensor[:-1].cumsum(0) - batched_edge_index = pyg_batch.edge_index.to(E.device) - total_nodes = pyg_batch.num_nodes # This is sum(num_nodes_list) + batched_edges = [] + for i, ei in enumerate(edge_index_list): + ei_device = ei.to(E.device) if ei.device != E.device else ei + batched_edges.append(ei_device + offsets[i]) + batched_edge_index = torch.cat(batched_edges, dim=1) + total_nodes = num_nodes_tensor.sum().item() - # 4. Run batched diffusion + # 4. Run batched diffusion + # Pass the entire batch of nodes and the single batched edge_index diffused_unpadded_E = _chebyshev_diffusion_batch( batched_edge_index, total_nodes, unpadded_E, k=k, - edge_weight=None, # Assuming edge_weight is None + edge_weight=None, # Assuming edge_weight is None as per GDTransformer's call beta=beta ) # 5. Re-pad the result + # Create a zero tensor with the original padded shape final_emb = torch.zeros_like(E) + + # Use the mask to "scatter" the diffused embeddings back to their + # original positions in the padded tensor final_emb[mask] = diffused_unpadded_E + assert final_emb.size() == E.size(), f"Expect {E.size()}, Got {final_emb.size()}" return final_emb \ No newline at end of file