fix: diffusion bottleneck

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

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
 
@@ -46,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_per_sample(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):
@@ -94,20 +122,50 @@ 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)
+    # 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)
 
-        # 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
+    # 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. 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_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
+    # 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 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

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,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"]).to("cuda")
-            gene_ids = batch["orig_gene_id"].detach().cpu().numpy()
-            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
+            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[gid] += emb_sum
+                gene_counts[gid] += int(cnt)
 
     # compute average embedding per gene
     gene_embeddings = {