Weekly GitHub Report for Xla

Thank you for subscribing to our weekly newsletter! Each week, we deliver a comprehensive summary of your GitHub project's latest activity right to your inbox, including an overview of your project's issues, pull requests, contributors, and commit activity.

Table of Contents

• I. News
  • 1.1. Recent Version Releases
  • 1.2. Other Noteworthy Updates
• II. Issues
  • 2.1. Top 5 Active Issues
  • 2.2. Top 5 Stale Issues
  • 2.3. Open Issues
  • 2.4. Closed Issues
  • 2.5. Issue Discussion Insights

• III. Pull Requests
  • 3.1. Open Pull Requests
  • 3.2. Closed Pull Requests
  • 3.3. Pull Request Discussion Insights

• IV. Contributors
  • 4.1. Contributors

I. News

1.1 Recent Version Releases:

No recent version releases were found.

1.2 Version Information:

Please provide the version release information you would like me to analyze and summarize.

II. Issues

2.1 Top 5 Active Issues:

We consider active issues to be issues that that have been commented on most frequently within the last week. Bot comments are omitted.

1. [BUG] [GOOGLE-TPU] [ERR: RUNTIME] [TPU] Silent BF16 gather corruption in T16-to-T8 fusion with debug callback (libtpu 0.0.48): This issue reports a silent data corruption bug occurring on TPU7x devices when running a specific JAX program that performs a BF16 gather fused with a T16-to-T8 layout conversion in the presence of a debug callback and FP8 quantization consumer. The corruption manifests as a precise block permutation of the gathered data, not a numerical tolerance error, and is reproducible only when the diagnostic callback is enabled or under certain memory layout conditions, implicating a defect in the TPU compiler’s handling of fused gather and layout conversion operations.

   • The comments provide extensive follow-up analysis confirming that the bug also occurs without the callback during full training runs, directly capturing corrupted BF16 data matching the same permutation; they detail shape- and memory-dependent triggering conditions, isolate the failure to a specific fused gather fusion with T16-to-T8 conversion, and request TPU compiler investigation and mitigation strategies while ruling out some potential causes and confirming the defect’s reproducibility and impact on training correctness.
   • Number of comments this week: 3

2. [GOOGLE-TPU] [STAT:AWAITING RESPONSE FROM CONTRIBUTOR] [ERR:PERFORMANCE] TPU v5p: async all-to-all and all-gather never overlap independent compute (reduce-scatter / all-reduce do, via SparseCore offload): This issue reports that on TPU v5p hardware, asynchronous all-to-all and all-gather collective operations inside shard_map do not overlap with independent matrix multiplication compute, unlike reduce-scatter and all-reduce which do overlap when offloaded to SparseCores. The user seeks guidance on whether there is a supported method or scheduler option to enable such overlap for all-to-all or all-gather collectives, as the current scheduling places the collective's completion immediately after its start, preventing latency hiding with concurrent compute.

   • The comments express appreciation for any improvements on compute/communication overlap and suggest trying specific environment variables to enable 2D all-gather offload and aggressive scheduling for all-to-all operations as potential workarounds.
   • Number of comments this week: 2

3. [BUG] [GPU] [ERR:PERFORMANCE] [XLA:GPU] Loop-fusion unroll factor 4 makes large complex128 fusions 1.4–3.6× slower than unroll 1 on A100 (sm_80); still 1.3–1.8× on main: This issue reports that on NVIDIA A100 GPUs (sm_80), the default loop fusion unroll factor of 4 in XLA causes large complex128 fusions to run significantly slower—up to 3.6× slower—than using an unroll factor of 1, due to increased register pressure and inability to vectorize wider loads. The problem persists even with the newer nvJitLink compiler and affects multi-output loop fusions such as adaptive ODE solver steps, with no current flag to disable or control the unroll factor on sm_80, leading to requests for capping unroll factors based on element size or making the factor configurable across architectures.

   • The single comment confirms the issue affects multi-output complex128 selects fused into one loop fusion, showing detailed benchmarks where unroll factor 4 causes 1.27–1.42× slowdown compared to separate kernels, while unroll 1 eliminates the penalty; it also analyzes PTX code behavior and suggests fixes including capping unroll factor or improving aliasing metadata to enable better load/store scheduling.
   • Number of comments this week: 1

4. [BUG] [GOOGLE-TPU] [XLA:TPU] Wrong results from conv_general_dilated with lhs_dilation=3: This issue reports that on TPU devices, the function lax.conv_general_dilated with lhs_dilation=3 produces incorrect convolution results without any error, diverging significantly from CPU and NumPy float64 reference outputs. The problem is reproducible across TPU versions and JAX releases, specifically affecting stride-3 transposed 1-D convolutions, and the user has provided detailed error measurements and system information to aid in diagnosis.

   • The single comment acknowledges the issue as a confirmed bug and indicates that a fix is being developed, with a promise to update once the fix is available.
   • Number of comments this week: 1

5. [BUG] [CPU] [NVIDIA-GPU] [GPU] [GOOGLE-TPU] [ERR: RUNTIME] [XLA] Wrong results when buffer assignment aliases a fusion's output with an operand it also reads through a bitcast: This issue describes a bug in the XLA compiler where buffer assignment incorrectly allows a fusion operation's output to alias the same buffer as one of its operands accessed through a bitcast, causing silent data corruption on H100 GPUs and CPUs. The problem arises because the buffer assignment checks only some positions of the operand for safe sharing, leading to overwritten elements before they are read, and the reporter provides a candidate patch that enforces stricter checks and fixes the issue on CPU.

   • The single comment acknowledges the report and confirms the ability to reproduce the bug, expressing intent to review and merge the proposed fix.
   • Number of comments this week: 1

2.2 Top 5 Stale Issues:

We consider stale issues to be issues that has had no activity within the last 30 days. The team should work together to get these issues resolved and closed as soon as possible.

As of our latest update, there are no stale issues for the project this week.

2.3 Open Issues

This section lists, groups, and then summarizes issues that were created within the last week in the repository.

Issues Opened This Week: 14

Summarized Issues:

• Optimization Limitations in All-Gather and Scatter Operations: The AllGatherOptimizer only optimizes addition operations involving all-gathered tensors but fails to optimize subtraction due to handling only commutative operations, leaving subtraction unoptimized. Similarly, the tpu_scatter_gather_optimizer pass rewrites scatter-add into sorted scatter causing TensorCore halts, indicating issues with optimization passes on TPU devices.
• issues/51155, issues/51393

• GPU and TPU Compilation and Performance Issues: Several issues report significant compilation or runtime problems on GPUs and TPUs, including drastic compilation time increases with GPU fusion autotuner on NVIDIA H100, a failure in GPU HloRematerialization causing compilation aborts, and a buffer assignment bug causing incorrect results on H100 GPUs, CPUs, and TPUs. These highlight challenges in compiler optimizations and memory management across hardware.
• issues/51338, issues/51432, issues/51490

• Incorrect Computation Results on TPU Devices: TPU-specific bugs cause incorrect results in various operations, such as BF16 gather output corruption with FP8 quantization, wrong convolution results with lhs_dilation=3 in lax.conv_general_dilated, and ragged_dot failing to compile with highest precision due to memory errors. These issues indicate hardware-specific computation inaccuracies and memory constraints.
• issues/51205, issues/51365, issues/51481

• Collective Operation Overlap and Communication Inefficiencies: On TPU v5p, asynchronous all-to-all and all-gather operations inside shard_map do not overlap with independent compute, unlike reduce-scatter and all-reduce, limiting performance. Additionally, the AllReduceSimplifier fails to eliminate product all-reduce on two replicas, causing unnecessary communication overhead.
• issues/51400, issues/51565

• Partitioning and Sharding Failures in XLA: The XLA SPMD partitioner crashes when handling empty tuple-shaped custom-calls due to improper sharding propagation, and the GpuAllReduceCombiner prematurely stops combining all-reduces after a large operation, missing opportunities for further optimization. These issues reflect robustness problems in partitioning and combination logic.
• issues/51433, issues/51444

• Loop Optimization Deficiencies: Duplicate scatter computations inside while loops are not eliminated by common subexpression elimination, resulting in redundant kernels and buffers, whereas identical computations outside loops are correctly deduplicated. This indicates a gap in loop-aware optimization passes.
• issues/51576

• Incorrect Layout and Data Movement Optimizations: A bug in JAX 0.12.0 on NVIDIA GH200 GPUs causes the triton-gemm-rewriter to incorrectly reorder layout-changing copies below bitcasts, leading to wrong double-precision dot product results without errors. This points to subtle issues in layout and data movement optimizations.
• issues/51186

2.4 Closed Issues

This section lists, groups, and then summarizes issues that were closed within the last week in the repository. This section also links the associated pull requests if applicable.

Issues Closed This Week: 5

Summarized Issues:

• GPU-specific crashes and compilation issues: Several issues highlight problems unique to GPU platforms, including a crash caused by a missing loop preheader terminator in StableHLO modules and a regression where the GemmFusionV2 pass incorrectly fuses all-to-all collectives into __triton_gemm kernels, leading to autotuner failures and compilation errors. These problems disrupt GPU execution and training workflows, particularly in multi-GPU expert-parallel MoE setups.
• [issues/47372, issues/51073]

• Collective operation bugs and performance on CPU: There are multiple issues related to collective operations on CPU, such as a performance bottleneck where all-reduce operations on tuples run one gloo ring per buffer causing high latency, and a correctness bug where lax.ppermute inside loops with multiple processes returns incorrect values due to premature collective starts and message tag conflicts. Both issues affect the efficiency and correctness of distributed computations on CPU platforms.
• [issues/50283, issues/50882]

• Compiler optimization passes causing long compile times: The PriorityFusion pass in the XLA compiler causes excessively long compile times that scale quadratically with model depth on GPU, due to duplication of residual normalization operations. Introducing optimization barriers can significantly reduce compile time and improve runtime performance, indicating a need for better handling of fusion passes in deep models.
• [issues/51256]

2.5 Issue Discussion Insights

This section will analyze the tone and sentiment of discussions within this project's open and closed issues that occurred within the past week. It aims to identify potentially heated exchanges and to maintain a constructive project environment.

Based on our analysis, there are no instances of toxic discussions in the project's open or closed issues from the past week.

III. Pull Requests

3.1 Open Pull Requests

This section provides a summary of pull requests that were opened in the repository over the past week. The top three pull requests with the highest number of commits are highlighted as 'key' pull requests. Other pull requests are grouped based on similar characteristics for easier analysis. Up to 25 pull requests are displayed in this section, while any remaining pull requests beyond this limit are omitted for brevity.

Pull Requests Opened This Week: 27

Key Open Pull Requests

1. [XLA] Memory-cost-aware latency hiding scheduler (kMemoryCost, aging, priced overlap): This pull request introduces a default-off, memory-cost-aware latency hiding scheduler for XLA that enhances scheduling decisions by incorporating new configuration knobs and rules—such as kMemoryCost, aging via kMemoryHeld, and an optional memory-priced overlap extension—to prioritize candidates that release large memory buffers more efficiently under high memory pressure, thereby aiming to reduce peak memory usage and improve scheduling for large training models without altering existing behavior when disabled.

• URL: pull/51322

• Associated Commits: 2ebcf, 32b77, 03e07, 29302, 65b73, 456d1, 085eb, 3ca76

2. [XLA:GPU] Fall back when memory fencing increases final allocations: This pull request introduces a fallback mechanism in the XLA GPU scheduler that compares fenced and unfenced memory scheduling candidates to prevent increases in final memory allocations caused by memory fencing, ensuring that the unfenced candidate is retained if fencing leads to larger temporary or total allocations, thereby safeguarding against regressions in memory usage without changing the default fencing behavior.

• URL: pull/51240

• Associated Commits: 5de90, 9ccc9, 7cb56, a3d64

3. Add intel iGPU definitions to oneapi_compute_capability: This pull request adds definitions for several new Intel integrated GPUs, including Meteor Lake, Arrow Lake, Lunar Lake, Panther Lake, Wildcat Lake, and Nova Lake variants, to the oneapi_compute_capability module to ensure proper recognition and construction of OneAPIComputeCapability objects for these devices.

• URL: pull/51168

• Associated Commits: f0a9d, 10f69, c6bee

Other Open Pull Requests

• ROCm version updates and test adjustments: Multiple pull requests update the ROCm backend to use ROCm 10.1.0 or 10.0 development tarballs, bumping the ROCm version in build configurations and skipping or disabling failing convolution tests due to known MIOpen bugs. These changes ensure compatibility with remote build execution environments and remove previous workarounds for blit fixes.
• pull/51226, pull/51451, pull/51536

• XLA GPU scheduler memory fencing improvements: Several pull requests enhance the XLA GPU scheduler by introducing fallback mechanisms that compare fenced and unfenced schedules, autotuning features for dynamic memory fencing thresholds, and an opt-in producer-side memory fencing pass. These improvements optimize memory usage by evaluating scheduling candidates and controlling buffer scheduling to reduce peak memory without regressing performance.
• pull/51427, pull/51404, pull/51301

• GPU memory allocator and buffer management fixes: Updates to the SKIP_PROFILED GPU memory allocator fix multi-GPU hangs on ROCm by allocating temporary buffers at reserved virtual addresses, and a bug fix ensures proper release of reservation-backed temporary buffers on execution setup failure. These changes reduce redundant buffer copies, prevent memory leaks, and improve stability across CUDA and ROCm platforms.
• pull/51248, pull/51247

• Atomic operation optimization in XLA GPU backend: A pull request improves the GPU backend by lowering atomic subtraction operations to atomic addition with negation, replacing costly compare-and-swap loops with efficient direct atomic instructions. This fix addresses performance regressions and includes extensive testing across multiple GPU architectures.
• pull/51488

• BlasLt mutex bug fix for hipBLASLt autotuning: A bug fix releases the BlasLt mutex before hipBLASLt algorithm queries to prevent blocking during autotuning, ensuring accurate timing and avoiding kernel timeouts on MI300X devices.
• pull/51236

• OneDNN custom-call AOT compilation support: This pull request adds Ahead Of Time compilation support for OneDNN custom-call thunks by enabling their serialization and deserialization, accompanied by extensive unit tests to validate the implementation.
• pull/51258

• Scaled-dot composites preservation on ROCm: A pull request ensures scaled-dot composites remain native during JAX custom partitioning on ROCm by preserving composite attributes and re-running the CompositeRewriter after SPMD partitioning. This enables multi-GPU compilation of GEMMs to native scaled-dot calls without changing single-device behavior.
• pull/51351

• OneDNN LayerNorm variance pattern update: Support is added for the updated OneDNN LayerNorm variance calculation pattern introduced by TensorFlow changes, restoring and improving performance by enabling fusion of operations that previously caused significant slowdowns.
• pull/51482

• Exposure of rocm_info struct for build tooling: The rocm_info struct is exposed from the local_config_rocm repository to external consumers, allowing build tools within the TensorFlow build root to locate required ROCm libraries by specifying rpaths.
• pull/51514

• Support for complex and complex core types: This pull request adds support for C32 and BC32 core types to XLA core HLO and Literals, including parsing, serialization, and evaluator support, laying the foundation for future complex f16/bf16 workloads while deferring backend numerical lowering.
• pull/51541

• MORI collectives backend enhancements: Support is added for contiguous sub-communicators within a single node, the MoriCommunicator is refactored to own its CollectivesFacade instance with rank-to-PE mapping validation, and the MORI dependency is repinned to a newer tag with matching facade-side support.
• pull/51563

• Algebraic simplifier bug fix for boolean select: A bug fix preserves the original result shape when simplifying boolean select operations with scalar predicates by broadcasting the predicate to match the tensor shape, preventing fatal shape compatibility errors and including regression tests.
• pull/51567

• GPU latency-hiding scheduler rule for memory reduction: A flag-gated scheduling rule prioritizes compute instructions that release large live buffers during asynchronous operations, reducing memory usage without stalls or exceeding limits, thus improving performance by shortening lifetimes of large intermediate buffers in training graphs.
• pull/51588

• oneDNN CPU batched matmul data race fix: A fix restricts K-partitioning reduction to batch size one for large contraction dimensions to prevent concurrent threads from writing to the same reduction buffer, avoiding non-deterministic outputs and severe numerical errors.
• pull/51589

• ROCm CI pipeline fix for argument order: The ROCm continuous integration pipeline is fixed by correcting the invalid order of arguments that caused compilation failures and remote build execution issues across multiple architectures.
• pull/51626

3.2 Closed Pull Requests

This section provides a summary of pull requests that were closed in the repository over the past week. The top three pull requests with the highest number of commits are highlighted as 'key' pull requests. Other pull requests are grouped based on similar characteristics for easier analysis. Up to 25 pull requests are displayed in this section, while any remaining pull requests beyond this limit are omitted for brevity.

Pull Requests Closed This Week: 40

Key Closed Pull Requests

1. [XLA] Fix Bazel 9 compatibility: This pull request aims to add compatibility for Bazel 9 to the XLA project by updating dependencies, explicitly loading external rules and providers, patching various third-party libraries and configurations, and ensuring that support for Bazel 8 is retained alongside these changes.

• URL: pull/51244

• Associated Commits: daa3e, d16d1, a02b6, 94014, 8372d, ab26c, 9e112, 64284, d0d70, 39a53, 2b78a, b9a18, fb6a6, 5fbe0, 22158, 97516, 0e354, f1c8b, bd3d2, 928a2

2. [XLA:GPU][oneAPI] Implement batch-pointers kernel for oneAPI backend: This pull request proposes adding a device-side batch-pointers kernel for the oneAPI backend that computes an array of batch pointers on the GPU using striding from a base address, registers this kernel as a free-function SYCL kernel identified by a sycl::kernel_id, and updates kernel loading in sycl_executor.cc to handle in-process symbol specifications via the SYCL host runtime, without introducing new tests since existing ones already cover this functionality.

• URL: pull/48684

• Associated Commits: 6c3ef, e9f36, b1427, 7466f, 25c82, a383f, 6b1ed, ece21, ef95c, 27e34, 8b457

3. Reduce XLA launch overhead by skipping unnecessary Prepare/Initialize for command buffers: This pull request proposes a performance optimization for XLA on GPU by introducing a mechanism to skip redundant Prepare and Initialize steps for command buffer thunks when their setup depends only on stable scopes like device or collective topology, thereby significantly reducing launch overhead—especially demonstrated by cutting nearly 800-900 microseconds from the preparation phase in large models such as DeepSeek 4.1-Flash on CUDA.

• URL: pull/50922

• Associated Commits: 0c8c6, 3e441, 2ee52, 3aed6, 09b2b, 3ea32, 970b4, 8e58d, d8b7a, b04be, 7da39

Other Closed Pull Requests

3.3 Pull Request Discussion Insights

This section will analyze the tone and sentiment of discussions within this project's open and closed pull requests that occurred within the past week. It aims to identify potentially heated exchanges and to maintain a constructive project environment.

Based on our analysis, there are no instances of toxic discussions in the project's open or closed pull requests from the past week.

IV. Contributors

4.1 Contributors

Active Contributors:

We consider an active contributor in this project to be any contributor who has made at least 1 commit, opened at least 1 issue, created at least 1 pull request, or made more than 2 comments in the last month.

If there are more than 10 active contributors, the list is truncated to the top 10 based on contribution metrics for better clarity.